feat(compute_ctl): pass compute type to pageserver with pg_options

Signed-off-by: Alex Chi Z <[email protected]>

Author: skyzh

compute_tools/src/compute.rs

@@ -855,6 +855,14 @@ impl ComputeNode {
             info!("Storage auth token not set");
         }
 
+        config.application_name("compute_ctl");
+        if let Some(spec) = &compute_state.pspec {
+            config.options(&format!(
+                "-c neon.endpoint_type={}",
+                spec.spec.mode.to_type_str()
+            ));
+        }
+
         // Connect to pageserver
         let mut client = config.connect(NoTls)?;
         let pageserver_connect_micros = start_time.elapsed().as_micros() as u64;

compute_tools/src/config.rs

@@ -99,6 +99,7 @@ pub fn write_postgres_conf(
         writeln!(file, "lc_numeric='C.UTF-8'")?;
     }
 
+    writeln!(file, "neon.endpoint_type={}", spec.mode.to_type_str())?;
     match spec.mode {
         ComputeMode::Primary => {}
         ComputeMode::Static(lsn) => {

libs/compute_api/src/spec.rs

@@ -272,6 +272,18 @@ pub enum ComputeMode {
     Replica,
 }
 
+impl ComputeMode {
+    /// Convert the compute mode to a string that can be used to identify the type of compute,
+    /// which means that if it's a static compute, the LSN will not be included.
+    pub fn to_type_str(&self) -> &'static str {
+        match self {
+            ComputeMode::Primary => "primary",
+            ComputeMode::Static(_) => "static",
+            ComputeMode::Replica => "replica",
+        }
+    }
+}
+
 /// Log level for audit logging
 /// Disabled, log, hipaa
 /// Default is Disabled

pageserver/src/page_service.rs

@@ -237,7 +237,7 @@ pub async fn libpq_listener_main(
 
 type ConnectionHandlerResult = anyhow::Result<()>;
 
-#[instrument(skip_all, fields(peer_addr, application_name))]
+#[instrument(skip_all, fields(peer_addr, application_name, endpoint_type))]
 #[allow(clippy::too_many_arguments)]
 async fn page_service_conn_main(
     conf: &'static PageServerConf,
@@ -2549,6 +2549,40 @@ where
             if let Some(app_name) = params.get("application_name") {
                 Span::current().record("application_name", field::display(app_name));
             }
+            if let Some(options) = params.get("options") {
+                let mut parsing_config = false;
+                for item in options.split_whitespace() {
+                    if item == "-c" {
+                        if !parsing_config {
+                            parsing_config = true;
+                        } else {
+                            // "-c" followed with another "-c"
+                            tracing::warn!("failed to parse the startup options: {options}");
+                            break;
+                        }
+                    } else if parsing_config {
+                        let Some((key, value)) = item.split_once('=') else {
+                            // "-c" followed with an invalid option
+                            tracing::warn!("failed to parse the startup options: {options}");
+                            break;
+                        };
+                        if key == "neon.endpoint_type" {
+                            Span::current().record("endpoint_type", field::display(value));
+                        } else {
+                            tracing::warn!("failed to parse the startup options: {options}");
+                            break;
+                        }
+                        parsing_config = false;
+                    } else {
+                        tracing::warn!("failed to parse the startup options: {options}");
+                        break;
+                    }
+                }
+                if parsing_config {
+                    // "-c" without the option
+                    tracing::warn!("failed to parse the startup options: {options}");
+                }
+            }
         };
 
         Ok(())
@@ -2662,6 +2696,7 @@ where
             PageServiceCmd::Set => {
                 // important because psycopg2 executes "SET datestyle TO 'ISO'"
                 // on connect
+                // TODO: allow setting options, i.e., application_name/endpoint_type via SET commands
                 pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
             }
             PageServiceCmd::LeaseLsn(LeaseLsnCmd {

pgxn/neon/libpagestore.c

@@ -50,6 +50,20 @@
 #define MIN_RECONNECT_INTERVAL_USEC 1000
 #define MAX_RECONNECT_INTERVAL_USEC 1000000
 
+
+enum NeonEndpointType {
+	EP_TYPE_PRIMARY = 0,
+	EP_TYPE_REPLICA,
+	EP_TYPE_STATIC
+};
+
+static const struct config_enum_entry neon_endpoint_types[] = {
+	{"primary", EP_TYPE_PRIMARY, false},
+	{"replica", EP_TYPE_REPLICA, false},
+	{"static", EP_TYPE_STATIC, false},
+	{NULL, 0, false}
+};
+
 /* GUCs */
 char	   *neon_timeline;
 char	   *neon_tenant;
@@ -62,6 +76,8 @@ int			flush_every_n_requests = 8;
 
 int         neon_protocol_version = 2;
 
+static int neon_endpoint_type = 0;
+
 static int	max_reconnect_attempts = 60;
 static int	stripe_size;
 
@@ -390,9 +406,10 @@ pageserver_connect(shardno_t shard_no, int elevel)
 	{
 	case PS_Disconnected:
 	{
-		const char *keywords[4];
-		const char *values[4];
+		const char *keywords[5];
+		const char *values[5];
 		char pid_str[16];
+		char endpoint_str[36];
 		int			n_pgsql_params;
 		TimestampTz	now;
 		int64		us_since_last_attempt;
@@ -464,6 +481,34 @@ pageserver_connect(shardno_t shard_no, int elevel)
 			n_pgsql_params++;
 		}
 
+		{
+			int ret = 0;
+			bool param_set = false;
+			switch (neon_endpoint_type)
+			{
+				case EP_TYPE_PRIMARY:
+					ret = snprintf(endpoint_str, sizeof(endpoint_str), "-c neon.endpoint_type=primary");
+					param_set = true;
+					break;
+				case EP_TYPE_REPLICA:
+					ret = snprintf(endpoint_str, sizeof(endpoint_str), "-c neon.endpoint_type=replica");
+					param_set = true;
+					break;
+				case EP_TYPE_STATIC:
+					ret = snprintf(endpoint_str, sizeof(endpoint_str), "-c neon.endpoint_type=static");
+					param_set = true;
+					break;
+			}
+			if (ret < 0 || ret >= (int)(sizeof(endpoint_str)))
+				elog(FATAL, "stack-allocated buffer too small to hold endpoint type");
+			if (param_set)
+			{
+				keywords[n_pgsql_params] = "options";
+				values[n_pgsql_params] = endpoint_str;
+				n_pgsql_params++;
+			}
+		}
+
 		keywords[n_pgsql_params] = NULL;
 		values[n_pgsql_params] = NULL;
 
@@ -1370,6 +1415,17 @@ pg_init_libpagestore(void)
 							GUC_UNIT_MS,
 							NULL, NULL, NULL);
 
+	DefineCustomEnumVariable(
+							"neon.endpoint_type",
+							"The compute endpoint node type",
+							NULL,
+							&neon_endpoint_type,
+							EP_TYPE_PRIMARY,
+							neon_endpoint_types,
+							PGC_POSTMASTER,
+							0,
+							NULL, NULL, NULL);
+
 	relsize_hash_init();
 
 	if (page_server != NULL)