Skip to content

Commit dc953de

Browse files
skyzhskyzh
authored
feat(pageserver): integrate PostHog with gc-compaction rollout (neondatabase#11917)
## Problem part of neondatabase#11813 ## Summary of changes * Integrate feature store with tenant structure. * gc-compaction picks up the current strategy from the feature store. * We only log them for now for testing purpose. They will not be used until we have more patches to support different strategies defined in PostHog. * We don't support property-based evaulation for now; it will be implemented later. * Evaluating result of the feature flag is not cached -- it's not efficient and cannot be used on hot path right now. * We don't report the evaluation result back to PostHog right now. I plan to enable it in staging once we get the patch merged. --------- Signed-off-by: Alex Chi Z <[email protected]>
1 parent 841517e commit dc953de

File tree

13 files changed

+267
-79
lines changed

13 files changed

+267
-79
lines changed

‎Cargo.lock‎

Lines changed: 6 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ azure_storage_blobs = { git = "https://github.com/neondatabase/azure-sdk-for-rus
247247
## Local libraries
248248
compute_api = { version = "0.1", path = "./libs/compute_api/" }
249249
consumption_metrics = { version = "0.1", path = "./libs/consumption_metrics/" }
250+
desim = { version = "0.1", path = "./libs/desim" }
250251
endpoint_storage = { version = "0.0.1", path = "./endpoint_storage/" }
251252
http-utils = { version = "0.1", path = "./libs/http-utils/" }
252253
metrics = { version = "0.1", path = "./libs/metrics/" }
@@ -259,19 +260,19 @@ postgres_backend = { version = "0.1", path = "./libs/postgres_backend/" }
259260
postgres_connection = { version = "0.1", path = "./libs/postgres_connection/" }
260261
postgres_ffi = { version = "0.1", path = "./libs/postgres_ffi/" }
261262
postgres_initdb = { path = "./libs/postgres_initdb" }
263+
posthog_client_lite = { version = "0.1", path = "./libs/posthog_client_lite" }
262264
pq_proto = { version = "0.1", path = "./libs/pq_proto/" }
263265
remote_storage = { version = "0.1", path = "./libs/remote_storage/" }
264266
safekeeper_api = { version = "0.1", path = "./libs/safekeeper_api" }
265267
safekeeper_client = { path = "./safekeeper/client" }
266-
desim = { version = "0.1", path = "./libs/desim" }
267268
storage_broker = { version = "0.1", path = "./storage_broker/" } # Note: main broker code is inside the binary crate, so linking with the library shouldn't be heavy.
268269
storage_controller_client = { path = "./storage_controller/client" }
269270
tenant_size_model = { version = "0.1", path = "./libs/tenant_size_model/" }
270271
tracing-utils = { version = "0.1", path = "./libs/tracing-utils/" }
271272
utils = { version = "0.1", path = "./libs/utils/" }
272273
vm_monitor = { version = "0.1", path = "./libs/vm_monitor/" }
273-
walproposer = { version = "0.1", path = "./libs/walproposer/" }
274274
wal_decoder = { version = "0.1", path = "./libs/wal_decoder" }
275+
walproposer = { version = "0.1", path = "./libs/walproposer/" }
275276

276277
## Common library dependency
277278
workspace_hack = { version = "0.1", path = "./workspace_hack/" }

‎libs/pageserver_api/src/config.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,21 @@ pub struct NodeMetadata {
4545
pub other: HashMap<String, serde_json::Value>,
4646
}
4747

48+
/// PostHog integration config.
49+
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
50+
pub struct PostHogConfig {
51+
/// PostHog project ID
52+
pub project_id: String,
53+
/// Server-side (private) API key
54+
pub server_api_key: String,
55+
/// Client-side (public) API key
56+
pub client_api_key: String,
57+
/// Private API URL
58+
pub private_api_url: String,
59+
/// Public API URL
60+
pub public_api_url: String,
61+
}
62+
4863
/// `pageserver.toml`
4964
///
5065
/// We use serde derive with `#[serde(default)]` to generate a deserializer
@@ -186,6 +201,8 @@ pub struct ConfigToml {
186201
pub tracing: Option<Tracing>,
187202
pub enable_tls_page_service_api: bool,
188203
pub dev_mode: bool,
204+
#[serde(skip_serializing_if = "Option::is_none")]
205+
pub posthog_config: Option<PostHogConfig>,
189206
pub timeline_import_config: TimelineImportConfig,
190207
#[serde(skip_serializing_if = "Option::is_none")]
191208
pub basebackup_cache_config: Option<BasebackupCacheConfig>,
@@ -701,6 +718,7 @@ impl Default for ConfigToml {
701718
import_job_checkpoint_threshold: NonZeroUsize::new(128).unwrap(),
702719
},
703720
basebackup_cache_config: None,
721+
posthog_config: None,
704722
}
705723
}
706724
}

‎libs/posthog_client_lite/Cargo.toml‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,14 @@ license.workspace = true
66

77
[dependencies]
88
anyhow.workspace = true
9+
arc-swap.workspace = true
910
reqwest.workspace = true
10-
serde.workspace = true
1111
serde_json.workspace = true
12+
serde.workspace = true
1213
sha2.workspace = true
13-
workspace_hack.workspace = true
1414
thiserror.workspace = true
15+
tokio = { workspace = true, features = ["process", "sync", "fs", "rt", "io-util", "time"] }
16+
tokio-util.workspace = true
17+
tracing-utils.workspace = true
18+
tracing.workspace = true
19+
workspace_hack.workspace = true

‎libs/posthog_client_lite/src/background_loop.rs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
//! A background loop that fetches feature flags from PostHog and updates the feature store.
2+
3+
use std::{sync::Arc, time::Duration};
4+
5+
use arc_swap::ArcSwap;
6+
use tokio_util::sync::CancellationToken;
7+
8+
use crate::{FeatureStore, PostHogClient, PostHogClientConfig};
9+
10+
/// A background loop that fetches feature flags from PostHog and updates the feature store.
11+
pub struct FeatureResolverBackgroundLoop {
12+
posthog_client: PostHogClient,
13+
feature_store: ArcSwap<FeatureStore>,
14+
cancel: CancellationToken,
15+
}
16+
17+
impl FeatureResolverBackgroundLoop {
18+
pub fn new(config: PostHogClientConfig, shutdown_pageserver: CancellationToken) -> Self {
19+
Self {
20+
posthog_client: PostHogClient::new(config),
21+
feature_store: ArcSwap::new(Arc::new(FeatureStore::new())),
22+
cancel: shutdown_pageserver,
23+
}
24+
}
25+
26+
pub fn spawn(self: Arc<Self>, handle: &tokio::runtime::Handle, refresh_period: Duration) {
27+
let this = self.clone();
28+
let cancel = self.cancel.clone();
29+
handle.spawn(async move {
30+
tracing::info!("Starting PostHog feature resolver");
31+
let mut ticker = tokio::time::interval(refresh_period);
32+
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
33+
loop {
34+
tokio::select! {
35+
_ = ticker.tick() => {}
36+
_ = cancel.cancelled() => break
37+
}
38+
let resp = match this
39+
.posthog_client
40+
.get_feature_flags_local_evaluation()
41+
.await
42+
{
43+
Ok(resp) => resp,
44+
Err(e) => {
45+
tracing::warn!("Cannot get feature flags: {}", e);
46+
continue;
47+
}
48+
};
49+
let feature_store = FeatureStore::new_with_flags(resp.flags);
50+
this.feature_store.store(Arc::new(feature_store));
51+
}
52+
tracing::info!("PostHog feature resolver stopped");
53+
});
54+
}
55+
56+
pub fn feature_store(&self) -> Arc<FeatureStore> {
57+
self.feature_store.load_full()
58+
}
59+
}

‎libs/posthog_client_lite/src/lib.rs‎

Lines changed: 38 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
//! A lite version of the PostHog client that only supports local evaluation of feature flags.
22
3+
mod background_loop;
4+
5+
pub use background_loop::FeatureResolverBackgroundLoop;
6+
37
use std::collections::HashMap;
48

59
use serde::{Deserialize, Serialize};
@@ -20,8 +24,7 @@ pub enum PostHogEvaluationError {
2024

2125
#[derive(Deserialize)]
2226
pub struct LocalEvaluationResponse {
23-
#[allow(dead_code)]
24-
flags: Vec<LocalEvaluationFlag>,
27+
pub flags: Vec<LocalEvaluationFlag>,
2528
}
2629

2730
#[derive(Deserialize)]
@@ -94,6 +97,12 @@ impl FeatureStore {
9497
}
9598
}
9699

100+
pub fn new_with_flags(flags: Vec<LocalEvaluationFlag>) -> Self {
101+
let mut store = Self::new();
102+
store.set_flags(flags);
103+
store
104+
}
105+
97106
pub fn set_flags(&mut self, flags: Vec<LocalEvaluationFlag>) {
98107
self.flags.clear();
99108
for flag in flags {
@@ -267,6 +276,7 @@ impl FeatureStore {
267276
&self,
268277
flag_key: &str,
269278
user_id: &str,
279+
properties: &HashMap<String, PostHogFlagFilterPropertyValue>,
270280
) -> Result<String, PostHogEvaluationError> {
271281
let hash_on_global_rollout_percentage =
272282
Self::consistent_hash(user_id, flag_key, "multivariate");
@@ -276,7 +286,7 @@ impl FeatureStore {
276286
flag_key,
277287
hash_on_global_rollout_percentage,
278288
hash_on_group_rollout_percentage,
279-
&HashMap::new(),
289+
properties,
280290
)
281291
}
282292

@@ -344,6 +354,19 @@ impl FeatureStore {
344354
}
345355
}
346356

357+
pub struct PostHogClientConfig {
358+
/// The server API key.
359+
pub server_api_key: String,
360+
/// The client API key.
361+
pub client_api_key: String,
362+
/// The project ID.
363+
pub project_id: String,
364+
/// The private API URL.
365+
pub private_api_url: String,
366+
/// The public API URL.
367+
pub public_api_url: String,
368+
}
369+
347370
/// A lite PostHog client.
348371
///
349372
/// At the point of writing this code, PostHog does not have a functional Rust client with feature flag support.
@@ -360,51 +383,30 @@ impl FeatureStore {
360383
/// want to report the feature flag usage back to PostHog. The current plan is to use PostHog only as an UI to
361384
/// configure feature flags so it is very likely that the client API will not be used.
362385
pub struct PostHogClient {
363-
/// The server API key.
364-
server_api_key: String,
365-
/// The client API key.
366-
client_api_key: String,
367-
/// The project ID.
368-
project_id: String,
369-
/// The private API URL.
370-
private_api_url: String,
371-
/// The public API URL.
372-
public_api_url: String,
386+
/// The config.
387+
config: PostHogClientConfig,
373388
/// The HTTP client.
374389
client: reqwest::Client,
375390
}
376391

377392
impl PostHogClient {
378-
pub fn new(
379-
server_api_key: String,
380-
client_api_key: String,
381-
project_id: String,
382-
private_api_url: String,
383-
public_api_url: String,
384-
) -> Self {
393+
pub fn new(config: PostHogClientConfig) -> Self {
385394
let client = reqwest::Client::new();
386-
Self {
387-
server_api_key,
388-
client_api_key,
389-
project_id,
390-
private_api_url,
391-
public_api_url,
392-
client,
393-
}
395+
Self { config, client }
394396
}
395397

396398
pub fn new_with_us_region(
397399
server_api_key: String,
398400
client_api_key: String,
399401
project_id: String,
400402
) -> Self {
401-
Self::new(
403+
Self::new(PostHogClientConfig {
402404
server_api_key,
403405
client_api_key,
404406
project_id,
405-
"https://us.posthog.com".to_string(),
406-
"https://us.i.posthog.com".to_string(),
407-
)
407+
private_api_url: "https://us.posthog.com".to_string(),
408+
public_api_url: "https://us.i.posthog.com".to_string(),
409+
})
408410
}
409411

410412
/// Fetch the feature flag specs from the server.
@@ -422,12 +424,12 @@ impl PostHogClient {
422424
// with bearer token of self.server_api_key
423425
let url = format!(
424426
"{}/api/projects/{}/feature_flags/local_evaluation",
425-
self.private_api_url, self.project_id
427+
self.config.private_api_url, self.config.project_id
426428
);
427429
let response = self
428430
.client
429431
.get(url)
430-
.bearer_auth(&self.server_api_key)
432+
.bearer_auth(&self.config.server_api_key)
431433
.send()
432434
.await?;
433435
let body = response.text().await?;
@@ -446,11 +448,11 @@ impl PostHogClient {
446448
) -> anyhow::Result<()> {
447449
// PUBLIC_URL/capture/
448450
// with bearer token of self.client_api_key
449-
let url = format!("{}/capture/", self.public_api_url);
451+
let url = format!("{}/capture/", self.config.public_api_url);
450452
self.client
451453
.post(url)
452454
.body(serde_json::to_string(&json!({
453-
"api_key": self.client_api_key,
455+
"api_key": self.config.client_api_key,
454456
"distinct_id": distinct_id,
455457
"event": event,
456458
"properties": properties,

0 commit comments

Comments
 (0)