feat(pegboard): implement ws wf

docs/libraries/workflow/GOTCHAS.md

@@ -140,3 +140,12 @@ Be careful when writing your struct definitions.
 When force waking a sleeping workflow by setting `wake_immediate = true`, know that if the workflow is
 currently on a `sleep` step it will go back to sleep if it has not reached its `wake_deadline` yet. For all
 other steps, the workflow will continue normally (usually just go back to sleep).
+
+## Long-lived tasks in `ctx.join`
+
+When executing multiple long-lived activities in a `ctx.join` call using a tuple, remember that internally it
+uses `tokio::join!` and not `tokio::try_join`. This means it will wait until all items finish and does not
+short circuit when an `Err` is returned from any branch.
+
+So if you have an activity that errors immediately and another that takes a while to finish, the `ctx.join`
+call will wait until the long task is complete (or errors) before returning.

lib/chirp-workflow/core/src/ctx/workflow.rs

@@ -541,17 +541,29 @@ impl WorkflowCtx {
 		Ok(output)
 	}
 
-	/// Joins multiple executable actions (activities, closures) and awaits them simultaneously.
+	/// Joins multiple executable actions (activities, closures) and awaits them simultaneously. This does not
+	/// short circuit in the event of an error to make sure activity side effects are recorded.
 	pub async fn join<T: Executable>(&mut self, exec: T) -> GlobalResult<T::Output> {
 		exec.execute(self).await
 	}
 
+	/// Joins multiple executable actions (activities, closures) and awaits them simultaneously, short
+	/// circuiting in the event of an error.
+	///
+	/// **BEWARE**: You should almost **never** use `try_join` over `join`.
+	///
+	/// The only possible case for using this over `join` is:
+	/// - You have long running activities that are cancellable
+	pub async fn try_join<T: Executable>(&mut self, exec: T) -> GlobalResult<T::Output> {
+		exec.try_execute(self).await
+	}
+
 	/// Spawns a new thread to execute workflow steps in.
 	pub fn spawn<F, T: Send + 'static>(&mut self, f: F) -> tokio::task::JoinHandle<GlobalResult<T>>
 	where
 		F: for<'a> FnOnce(&'a mut WorkflowCtx) -> AsyncResult<'a, T> + Send + 'static,
 	{
-		let mut ctx = self.clone();
+		let mut ctx = self.branch();
 		tokio::task::spawn(async move { closure(f).execute(&mut ctx).await })
 	}
 

lib/chirp-workflow/core/src/executable.rs

@@ -8,10 +8,16 @@ use crate::ctx::WorkflowCtx;
 /// Signifies a retryable executable entity in a workflow. For example: activity, tuple of activities (join),
 /// closure.
 #[async_trait]
-pub trait Executable: Send {
+pub trait Executable: Send + Sized {
 	type Output: Send;
 
 	async fn execute(self, ctx: &mut WorkflowCtx) -> GlobalResult<Self::Output>;
+	/// In the event that an executable has multiple sub executables (i.e. a tuple of executables), this can
+	/// be implemented to provide the ability to choose whether to use `tokio::join!` or tokio::try_join!`
+	/// internally. Default implementation just calls `execute`.
+	async fn try_execute(self, ctx: &mut WorkflowCtx) -> GlobalResult<Self::Output> {
+		self.execute(ctx).await
+	}
 }
 
 pub type AsyncResult<'a, T> = Pin<Box<dyn Future<Output = GlobalResult<T>> + Send + 'a>>;
@@ -59,6 +65,23 @@ macro_rules! impl_tuple {
 				// Handle errors here instead
 				Ok(($($args?),*))
 			}
+
+			async fn try_execute(self, ctx: &mut WorkflowCtx) -> GlobalResult<Self::Output> {
+				#[allow(non_snake_case)]
+				let ($($args),*) = self;
+
+				#[allow(non_snake_case)]
+				let ($(mut $args),*) = ($(
+					TupleHelper {
+						branch: ctx.step(),
+						exec: $args,
+					}
+				),*);
+
+				tokio::try_join!(
+					$($args.exec.execute(&mut $args.branch)),*
+				)
+			}
 		}
 	}
 }

svc/Cargo.toml

@@ -178,6 +178,7 @@ members = [
 	"pkg/monolith/standalone/workflow-worker",
 	"pkg/nomad/standalone/monitor",
 	"pkg/nsfw/ops/image-score",
+	"pkg/pegboard",
 	"pkg/perf/ops/log-get",
 	"pkg/profanity/ops/check",
 	"pkg/region/ops/get",

svc/pkg/pegboard/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "pegboard"
+version = "0.0.1"
+edition = "2018"
+authors = ["Rivet Gaming, LLC <[email protected]>"]
+license = "Apache-2.0"
+
+[dependencies]
+chirp-workflow = { path = "../../../lib/chirp-workflow/core" }
+tokio-tungstenite = "0.23.1"
+serde = { version = "1.0.198", features = ["derive"] }

svc/pkg/pegboard/Service.toml

@@ -0,0 +1,7 @@
+[service]
+name = "pegboard"
+
+[runtime]
+kind = "rust"
+
+[package]

svc/pkg/pegboard/src/lib.rs

@@ -0,0 +1,14 @@
+use chirp_workflow::prelude::*;
+
+pub mod ops;
+pub mod utils;
+pub mod workflows;
+
+pub fn registry() -> WorkflowResult<Registry> {
+	use workflows::*;
+
+	let mut registry = Registry::new();
+	registry.register_workflow::<ws::Workflow>()?;
+
+	Ok(registry)
+}

svc/pkg/pegboard/src/ops/mod.rs

@@ -0,0 +1 @@
+

svc/pkg/pegboard/src/utils.rs

@@ -0,0 +1 @@
+

svc/pkg/pegboard/src/workflows/mod.rs

@@ -0,0 +1 @@
+pub mod ws;

svc/pkg/pegboard/src/workflows/ws.rs

@@ -0,0 +1,81 @@
+use std::{
+	collections::HashMap,
+	net::SocketAddr,
+	sync::{Arc, RwLock},
+};
+
+use chirp_workflow::prelude::*;
+use futures_util::FutureExt;
+// use tokio::net::{TcpListener, TcpStream};
+// use tokio_tungstenite::tungstenite::protocol::Message;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Input {}
+
+#[workflow]
+pub async fn pegboard_ws(ctx: &mut WorkflowCtx, input: &Input) -> GlobalResult<()> {
+	// let addr = "127.0.0.1:8080";
+	// let listener = TcpListener::bind(&addr).await?;
+	// println!("Listening on: {}", addr);
+
+	let conns = Arc::new(RwLock::new(HashMap::<(), ()>::new()));
+
+	ctx.try_join((
+		closure(|ctx| socket_thread(ctx, conns.clone()).boxed()),
+		closure(|ctx| signal_thread(ctx, conns.clone()).boxed()),
+	))
+	.await?;
+
+	Ok(())
+}
+
+async fn socket_thread(
+	ctx: &mut WorkflowCtx,
+	conns: Arc<RwLock<HashMap<(), ()>>>,
+) -> GlobalResult<()> {
+	ctx.repeat(|ctx| {
+		async move {
+			if let Ok((stream, addr)) = listener.accept().await {
+				handle_connection(stream, addr).await;
+			} else {
+				tracing::error!("failed to connect websocket");
+			}
+
+			Ok(Loop::Continue)
+		}.boxed()
+	)
+
+	Ok(())
+}
+
+async fn signal_thread(
+	ctx: &mut WorkflowCtx,
+	conns: Arc<RwLock<HashMap<(), ()>>>,
+) -> GlobalResult<()> {
+	Ok(())
+}
+
+async fn handle_connection(ctx, raw_stream: TcpStream, addr: SocketAddr) {
+	ctx.spawn(|ctx| async move {
+		let ws_stream = tokio_tungstenite::accept_async(raw_stream).await?;
+		let (mut write, mut read) = ws_stream.split();
+
+		println!("New WebSocket connection: {}", addr);
+
+		while let Some(Ok(msg)) = read.next().await {
+			if msg.is_text() || msg.is_binary() {
+				write.send(msg).await?;
+			}
+		}
+
+		Ok(())
+	}.boxed()).await
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+struct FooInput {}
+
+#[activity(Foo)]
+async fn foo(ctx: &ActivityCtx, input: &FooInput) -> GlobalResult<()> {
+	Ok(())
+}