fix(pb): use runner protocol for serverless (#3020)

Author: MasterPtato

Cargo.lock

@@ -1,7 +1,7 @@
 // Based off of https://github.com/tokio-rs/tracing-opentelemetry/blob/v0.1.x/examples/opentelemetry-otlp.rs
 
 use console_subscriber;
-use opentelemetry::trace::TracerProvider as _;
+use opentelemetry::trace::TracerProvider;
 use rivet_metrics::OtelProviderGuard;
 use tracing_opentelemetry::{MetricsLayer, OpenTelemetryLayer};
 use tracing_subscriber::{EnvFilter, Layer, layer::SubscriberExt, util::SubscriberInitExt};

packages/common/runtime/src/traces.rs

@@ -9,7 +9,7 @@ use std::{
 };
 use tokio::sync::{Mutex, mpsc};
 use universalpubsub::{NextOutput, PubSub, PublishOpts, Subscriber};
-use vbare::OwnedVersionedData as _;
+use vbare::OwnedVersionedData;
 
 const GC_INTERVAL: Duration = Duration::from_secs(60);
 const MESSAGE_ACK_TIMEOUT: Duration = Duration::from_secs(5);

packages/core/pegboard-gateway/src/shared_state.rs

@@ -9,7 +9,7 @@ use rivet_guard_core::websocket_handle::WebSocketReceiver;
 use rivet_runner_protocol::{self as protocol, PROTOCOL_VERSION, versioned};
 use std::sync::{Arc, atomic::Ordering};
 use universalpubsub::PublishOpts;
-use vbare::OwnedVersionedData as _;
+use vbare::OwnedVersionedData;
 
 use crate::{
 	conn::Conn,

packages/core/pegboard-runner/src/client_to_pubsub_task.rs

@@ -13,7 +13,7 @@ use std::{
 	time::Duration,
 };
 use tokio::sync::Mutex;
-use vbare::OwnedVersionedData as _;
+use vbare::OwnedVersionedData;
 
 use crate::{errors::WsError, utils::UrlData};
 

packages/core/pegboard-runner/src/conn.rs

@@ -5,7 +5,7 @@ use hyper_tungstenite::tungstenite::Message as WsMessage;
 use rivet_runner_protocol::{self as protocol, versioned};
 use std::sync::Arc;
 use universalpubsub::{NextOutput, Subscriber};
-use vbare::OwnedVersionedData as _;
+use vbare::OwnedVersionedData;
 
 use crate::{
 	conn::{Conn, TunnelActiveRequest},

packages/core/pegboard-runner/src/pubsub_to_client_task.rs

@@ -7,6 +7,7 @@ edition.workspace = true
 
 [dependencies]
 anyhow.workspace = true
+base64.workspace = true
 epoxy.workspace = true
 gas.workspace = true
 reqwest-eventsource.workspace = true
@@ -16,6 +17,7 @@ rivet-runner-protocol.workspace = true
 rivet-types.workspace = true
 tracing.workspace = true
 universaldb.workspace = true
+vbare.workspace = true
 
 namespace.workspace = true
 pegboard.workspace = true

packages/core/pegboard-serverless/Cargo.toml

@@ -7,6 +7,8 @@ use std::{
 };
 
 use anyhow::Result;
+use base64::Engine;
+use base64::engine::general_purpose::STANDARD as BASE64;
 use futures_util::{StreamExt, TryStreamExt};
 use gas::prelude::*;
 use pegboard::keys;
@@ -17,6 +19,7 @@ use rivet_types::namespaces::RunnerConfig;
 use tokio::{sync::oneshot, task::JoinHandle, time::Duration};
 use universaldb::options::StreamingMode;
 use universaldb::utils::IsolationLevel::*;
+use vbare::OwnedVersionedData;
 
 const X_RIVET_TOKEN: HeaderName = HeaderName::from_static("x-rivet-token");
 
@@ -247,7 +250,17 @@ async fn outbound_handler(
 					tracing::debug!(%msg.data, "received outbound req message");
 
 					if runner_id.is_none() {
-						runner_id = Some(Id::parse(&msg.data)?);
+						let data = BASE64.decode(msg.data).context("invalid base64 message")?;
+						let payload =
+							protocol::versioned::ToServerlessServer::deserialize_with_embedded_version(&data)
+								.context("invalid payload")?;
+
+						match payload {
+							protocol::ToServerlessServer::ToServerlessServerInit(init) => {
+								runner_id =
+									Some(Id::parse(&init.runner_id).context("invalid runner id")?);
+							}
+						}
 					}
 				}
 				Err(sse::Error::StreamEnded) => {

packages/core/pegboard-serverless/src/lib.rs

@@ -1,4 +1,4 @@
-use base64::Engine as _;
+use base64::Engine;
 use base64::prelude::BASE64_STANDARD;
 use futures_util::StreamExt;
 use futures_util::{FutureExt, TryStreamExt};

packages/services/pegboard/src/workflows/actor/runtime.rs

@@ -7,7 +7,7 @@ use universaldb::{
 	utils::{FormalChunkedKey, IsolationLevel::*},
 };
 use universalpubsub::PublishOpts;
-use vbare::OwnedVersionedData as _;
+use vbare::OwnedVersionedData;
 
 use crate::{keys, workflows::actor::Allocate};