fix(pegboard): convert ws manager socket to unix socket

MasterPtato · MasterPtato · commit bfa4fd3dda76 · 2025-06-17T18:38:33.000Z
diff --git a/docker/dev-full/prometheus/prometheus.yml b/docker/dev-full/prometheus/prometheus.yml
@@ -14,4 +14,4 @@ scrape_configs:
   - job_name: rivet-client
     static_configs:
       - targets:
-        - rivet-client:6090
+        - rivet-client:8091
diff --git a/docker/dev-full/vector-client/vector.yaml b/docker/dev-full/vector-client/vector.yaml
@@ -8,7 +8,7 @@ sources:
   prometheus_pegboard:
     type: prometheus_scrape
     endpoints:
-      - http://rivet-client:6090
+      - http://rivet-client:8091
     scrape_interval_secs: 15
   
   dynamic_events_http:
diff --git a/docker/monolith/build-scripts/setup_s6.ts b/docker/monolith/build-scripts/setup_s6.ts
@@ -103,7 +103,7 @@ const services: Service[] = [
 		rootUser: true,
 		ports: {
 			runner: 6080,
-			metrics: 6090,
+			metrics: 8091,
 		},
 	},
 
diff --git a/docker/monolith/vector-client/vector.yaml b/docker/monolith/vector-client/vector.yaml
@@ -12,7 +12,7 @@ sources:
   prometheus_pegboard:
     type: prometheus_scrape
     endpoints:
-      - http://rivet-client:6090
+      - http://rivet-client:8091
     scrape_interval_secs: 15
   
   pegboard_manager:
diff --git a/packages/common/chirp-workflow/core/src/worker.rs b/packages/common/chirp-workflow/core/src/worker.rs
@@ -83,6 +83,14 @@ impl Worker {
 		loop {
 			tokio::select! {
 				_ = tick_interval.tick() => {},
+				res = wake_sub.next() => {
+					if res.is_none() {
+						return Err(WorkflowError::SubscriptionUnsubscribed.into());
+					}
+
+					tick_interval.reset();
+				},
+
 				res = &mut gc_handle => {
 					tracing::error!(?res, "metrics task unexpectedly stopped");
 					break;
@@ -91,13 +99,6 @@ impl Worker {
 					tracing::error!(?res, "metrics task unexpectedly stopped");
 					break;
 				},
-				res = wake_sub.next() => {
-					if res.is_none() {
-						return Err(WorkflowError::SubscriptionUnsubscribed.into());
-					}
-
-					tick_interval.reset();
-				},
 				_ = ctrl_c() => break,
 				_ = sigterm.recv() => break,
 			}
diff --git a/packages/common/metrics/src/server.rs b/packages/common/metrics/src/server.rs
@@ -1,11 +1,12 @@
+use std::net::SocketAddr;
+
 use global_error::prelude::*;
 use hyper::{
 	header::CONTENT_TYPE,
 	service::{make_service_fn, service_fn},
 	Body, Request, Response, Server,
 };
 use prometheus::{Encoder, TextEncoder};
-use std::net::SocketAddr;
 
 // TODO: Record extra labels
 
diff --git a/packages/core/services/cluster/src/workflows/server/install/install_scripts/mod.rs b/packages/core/services/cluster/src/workflows/server/install/install_scripts/mod.rs
@@ -179,7 +179,7 @@ pub async fn gen_initialize(
 				"pegboard".into(),
 				components::vector::PrometheusTarget {
 					// Should match port from pb manager config
-					endpoint: "http://127.0.0.1:6090".into(),
+					endpoint: "http://127.0.0.1:8091".into(),
 					scrape_interval: 15,
 				},
 			);
@@ -194,7 +194,7 @@ pub async fn gen_initialize(
 				"pegboard".into(),
 				components::vector::PrometheusTarget {
 					// Should match port from pb manager config
-					endpoint: "http://127.0.0.1:6090".into(),
+					endpoint: "http://127.0.0.1:8091".into(),
 					scrape_interval: 15,
 				},
 			);
diff --git a/packages/edge/infra/client/config/Cargo.toml b/packages/edge/infra/client/config/Cargo.toml
@@ -6,10 +6,13 @@ license.workspace = true
 edition.workspace = true
 
 [dependencies]
+anyhow = "1.0"
 ipnet = { version = "2.10.1", features = ["serde"] }
+pegboard.workspace = true
+rivet-util.workspace = true
 schemars = { version = "0.8.21", features = ["url", "uuid1"] }
 serde = { version = "1.0.195", features = ["derive"] }
+serde_json = "1.0"
+tokio-util = { version = "0.7", features = ["codec"] }
 url = "2.2.2"
 uuid = { version = "1.6.1", features = ["v4"] }
-pegboard.workspace = true
-rivet-util.workspace = true
diff --git a/packages/edge/infra/client/config/src/manager.rs b/packages/edge/infra/client/config/src/manager.rs
@@ -94,9 +94,6 @@ pub struct Runner {
 	/// ````
 	pub use_resource_constraints: Option<bool>,
 
-	/// WebSocket Port for runners on this machine to connect to.
-	pub port: Option<u16>,
-
 	pub container_runner_binary_path: Option<PathBuf>,
 }
 
@@ -109,10 +106,6 @@ impl Runner {
 		self.use_resource_constraints.unwrap_or(true)
 	}
 
-	pub fn port(&self) -> u16 {
-		self.port.unwrap_or(6080)
-	}
-
 	pub fn container_runner_binary_path(&self) -> PathBuf {
 		self.container_runner_binary_path
 			.clone()
@@ -259,7 +252,7 @@ pub struct Metrics {
 
 impl Metrics {
 	pub fn port(&self) -> u16 {
-		self.port.unwrap_or(6090)
+		self.port.unwrap_or(8091)
 	}
 }
 
diff --git a/packages/edge/infra/client/config/src/runner_protocol.rs b/packages/edge/infra/client/config/src/runner_protocol.rs
@@ -1,5 +1,9 @@
+use anyhow::*;
 use pegboard::protocol;
 use serde::{Deserialize, Serialize};
+use tokio_util::codec::LengthDelimitedCodec;
+use serde::{de::DeserializeOwned, Serialize};
+use bytes::Bytes;
 
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case", deny_unknown_fields)]
@@ -38,3 +42,38 @@ pub enum ActorState {
 	Running,
 	Exited { exit_code: Option<i32> },
 }
+
+pub fn codec() -> LengthDelimitedCodec {
+	LengthDelimitedCodec::builder()
+		.length_field_type::<u32>()
+		.length_field_length(4)
+		// No offset
+		.length_field_offset(0)
+		// Header length is not included in the length calculation
+		.length_adjustment(4)
+		// Header is included in the returned bytes
+		.num_skip(0)
+		.new_codec()
+}
+
+pub fn encode_frame<T: Serialize>(payload: &T) -> Result<Vec<u8>> {
+	let mut buf = Vec::with_capacity(4);
+	buf.extend_from_slice(&[0u8; 4]); // header (currently unused)
+
+	let mut cursor = Cursor::new(&mut buf);
+	serde_json::to_writer(&mut cursor, payload)?;
+
+	Ok(buf)
+}
+
+fn decode_frame<T: DeserializeOwned>(frame: &Vec<u8>) -> Result<([u8; 4], T)> {
+	ensure!(frame.len() >= 4, "Frame too short");
+
+	// Extract the header (first 4 bytes)
+	let header = [frame[0], frame[1], frame[2], frame[3]];
+
+	// Deserialize the rest of the frame (payload after the header)
+	let payload = serde_json::from_slice(&frame[4..])?;
+
+	Ok((header, payload))
+}
diff --git a/packages/edge/infra/client/manager/src/ctx.rs b/packages/edge/infra/client/manager/src/ctx.rs
@@ -220,44 +220,6 @@ impl Ctx {
 
 		self.receive_init(&mut rx).await?;
 
-		// Start runner socket and attaches incoming connections to their corresponding runner
-		let self2 = self.clone();
-		let runner_socket: tokio::task::JoinHandle<Result<()>> = tokio::spawn(async move {
-			tracing::info!(port=%self2.config().runner.port(), "listening for runner sockets");
-
-			let listener =
-				TcpListener::bind((self2.config().runner.ip(), self2.config().runner.port()))
-					.await
-					.map_err(RuntimeError::RunnerSocketListenFailed)?;
-
-			loop {
-				match listener.accept().await {
-					Ok((stream, _)) => {
-						let mut ws_stream = Some(tokio_tungstenite::accept_async(stream).await?);
-
-						tracing::info!("received new socket");
-
-						if let Err(err) = self2.receive_runner_init_message(&mut ws_stream).await {
-							tracing::error!(
-								?err,
-								"failed to receive init message from runner socket"
-							);
-						}
-
-						// Close stream
-						if let Some(mut ws_stream) = ws_stream {
-							let close_frame = CloseFrame {
-								code: CloseCode::Error,
-								reason: "init failed".into(),
-							};
-							ws_stream.send(Message::Close(Some(close_frame))).await?;
-						}
-					}
-					Err(err) => tracing::error!(?err, "failed to connect websocket"),
-				}
-			}
-		});
-
 		// Start ping thread after init packet is received because ping denotes this client as "ready"
 		let self2 = self.clone();
 		let ping_thread: tokio::task::JoinHandle<Result<()>> = tokio::spawn(async move {
@@ -540,63 +502,6 @@ impl Ctx {
 		Ok(())
 	}
 
-	async fn receive_runner_init_message(
-		self: &Arc<Self>,
-		ws_stream: &mut Option<WebSocketStream<TcpStream>>,
-	) -> Result<()> {
-		match tokio::time::timeout(
-			RUNNER_INIT_TIMEOUT,
-			ws_stream.as_mut().context("ws_stream should exist")?.next(),
-		)
-		.await
-		{
-			Ok(msg) => match msg {
-				Some(Ok(Message::Close(_))) | None => {
-					tracing::debug!("runner socket closed");
-					Ok(())
-				}
-				Some(Ok(Message::Binary(buf))) => {
-					match serde_json::from_slice::<runner_protocol::ToManager>(&buf) {
-						Ok(runner_protocol::ToManager::Init { access_token }) => {
-							let claims = Claims::decode(&access_token, self.secret())
-								.context("could not decode access token from init packet")?;
-
-							// Read runner id from claims
-							let runner_id = 'ent: {
-								for ent in claims.ent() {
-									#[allow(irrefutable_let_patterns)]
-									if let Entitlement::Runner { runner_id } = ent {
-										break 'ent runner_id;
-									}
-								}
-
-								bail!("no runner entitlement");
-							};
-
-							if let Some(runner) = self.runners.read().await.get(&runner_id) {
-								runner
-									.attach_socket(
-										self,
-										ws_stream.take().context("ws_stream should exist")?,
-									)
-									.await?;
-
-								Ok(())
-							} else {
-								bail!("killing unknown runner: {runner_id}");
-							}
-						}
-						Ok(packet) => bail!("invalid init packet from runner: {packet:?}"),
-						Err(err) => Err(err).context("invalid init packet from runner"),
-					}
-				}
-				Some(Ok(packet)) => bail!("invalid init packet from runner: {packet:?}"),
-				Some(Err(err)) => Err(err).context("runner socket error"),
-			},
-			Err(_) => bail!("runner socket init timed out"),
-		}
-	}
-
 	/// Returns None if the runner could not be found in the runners map on time.
 	async fn get_or_create_runner(
 		&self,
@@ -909,6 +814,10 @@ impl Ctx {
 			let runner = runner.clone();
 			let self2 = self.clone();
 			tokio::spawn(async move {
+				// The socket file already exists, this will reconnect and spawn the appropriate task to
+				// handle the connection
+				runner.start_socket(&self2).await?;
+
 				if let Err(err) = runner.observe(&self2, false).await {
 					tracing::error!(?runner_id, ?err, "observe failed");
 				}
diff --git a/packages/edge/infra/client/manager/src/main.rs b/packages/edge/infra/client/manager/src/main.rs
@@ -176,7 +176,7 @@ async fn run(init: Init, first: bool) -> Result<()> {
 	}
 
 	// Start metrics server
-	let metrics_thread = tokio::spawn(metrics::run_standalone(init.config.client.metrics.port()));
+	let metrics_task = tokio::spawn(metrics::run_standalone(init.config.client.metrics.port()));
 
 	let url = build_ws_url(&init.config).await?;
 	tracing::info!("connecting to pegboard ws: {}", &url);
@@ -195,7 +195,7 @@ async fn run(init: Init, first: bool) -> Result<()> {
 	let ctx = Ctx::new(init.config, init.system, init.secret, init.pool, tx);
 
 	tokio::try_join!(
-		async { metrics_thread.await?.map_err(Into::into) },
+		metrics_task,
 		ctx.run(rx),
 	)?;
 
diff --git a/packages/edge/infra/client/manager/src/runner/mod.rs b/packages/edge/infra/client/manager/src/runner/mod.rs
diff --git a/packages/edge/infra/client/manager/src/runner/oci_config.rs b/packages/edge/infra/client/manager/src/runner/oci_config.rs
diff --git a/packages/edge/infra/client/manager/src/runner/setup.rs b/packages/edge/infra/client/manager/src/runner/setup.rs

Original file line number	Diff line number	Diff line change
`@@ -94,9 +94,6 @@ pub struct Runner {`
`94`	`94`	/// ````
`95`	`95`	`pub use_resource_constraints: Option<bool>,`
`96`	`96`
`97`		`- /// WebSocket Port for runners on this machine to connect to.`
`98`		`- pub port: Option<u16>,`
`99`		`-`
`100`	`97`	`pub container_runner_binary_path: Option<PathBuf>,`
`101`	`98`	`}`
`102`	`99`
`@@ -109,10 +106,6 @@ impl Runner {`
`109`	`106`	`self.use_resource_constraints.unwrap_or(true)`
`110`	`107`	`}`
`111`	`108`
`112`		`- pub fn port(&self) -> u16 {`
`113`		`- self.port.unwrap_or(6080)`
`114`		`- }`
`115`		`-`
`116`	`109`	`pub fn container_runner_binary_path(&self) -> PathBuf {`
`117`	`110`	`self.container_runner_binary_path`
`118`	`111`	`.clone()`
`@@ -259,7 +252,7 @@ pub struct Metrics {`
`259`	`252`
`260`	`253`	`impl Metrics {`
`261`	`254`	`pub fn port(&self) -> u16 {`
`262`		`- self.port.unwrap_or(6090)`
	`255`	`+ self.port.unwrap_or(8091)`
`263`	`256`	`}`
`264`	`257`	`}`
`265`	`258`