merge: #3470

3470: feat(si-layer-cache): impl rebaser requests work queue stream r=adamhjk a=fnichol

This is an initial implementation of a work queue stream of rebaser requests, suitable for use by the rebaser service.

Note that it's not yet clear if additional work queue streams follow this pattern, if the code is extractable into the service-related crates, or if some other code de-duplication is possible. At the moment, the newly added Rust types are doing a *lot* of heavy lifting to ensure correctness of behavior and type safety.

Also note that this NATS Jetstream setup may require further testing and refinement to confirm the desired stream sourcing behavior.

<img src="https://media1.giphy.com/media/5YuhLwDgrgtRVwI7OY/giphy.gif"/>

Co-authored-by: Fletcher Nichol <fletcher@systeminit.com>

Author: si-bors-ng[bot]

lib/si-layer-cache/src/activities.rs

@@ -283,3 +283,107 @@ impl Stream for ActivityStream {
         }
     }
 }
+
+#[derive(Clone)]
+pub struct AckRebaseRequest {
+    pub id: ActivityId,
+    pub payload: RebaseRequest,
+    pub metadata: LayeredEventMetadata,
+    acker: Arc<Acker>,
+}
+
+impl AckRebaseRequest {
+    pub async fn ack(&self) -> LayerDbResult<()> {
+        self.acker.ack().await.map_err(LayerDbError::NatsAck)
+    }
+
+    pub async fn ack_with(&self, kind: AckKind) -> LayerDbResult<()> {
+        self.acker
+            .ack_with(kind)
+            .await
+            .map_err(LayerDbError::NatsAck)
+    }
+
+    pub async fn double_ack(&self) -> LayerDbResult<()> {
+        self.acker.double_ack().await.map_err(LayerDbError::NatsAck)
+    }
+}
+
+pub struct RebaserRequestsWorkQueueStream {
+    inner: jetstream::consumer::pull::Stream,
+}
+
+impl RebaserRequestsWorkQueueStream {
+    const CONSUMER_NAME: &'static str = "rebaser-requests";
+
+    pub(crate) async fn create(nats_client: &NatsClient) -> LayerDbResult<Self> {
+        let context = jetstream::new(nats_client.as_inner().clone());
+
+        // Ensure the sourced stream is created
+        let _activities =
+            nats::layerdb_activities_stream(&context, nats_client.metadata().subject_prefix())
+                .await?;
+
+        let inner = nats::rebaser_requests_work_queue_stream(
+            &context,
+            nats_client.metadata().subject_prefix(),
+        )
+        .await?
+        .create_consumer(Self::consumer_config())
+        .await?
+        .messages()
+        .await?;
+
+        Ok(Self { inner })
+    }
+
+    #[inline]
+    fn consumer_config() -> jetstream::consumer::pull::Config {
+        jetstream::consumer::pull::Config {
+            durable_name: Some(Self::CONSUMER_NAME.to_string()),
+            description: Some("rebaser requests consumer".to_string()),
+            ..Default::default()
+        }
+    }
+}
+
+impl Stream for RebaserRequestsWorkQueueStream {
+    type Item = LayerDbResult<AckRebaseRequest>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match Pin::new(&mut self.inner.next()).poll(cx) {
+            // Process the message
+            Poll::Ready(Some(Ok(msg))) => {
+                let (msg, acker) = msg.split();
+
+                match postcard::from_bytes::<Activity>(&msg.payload) {
+                    // Successfully deserialized into an activity
+                    Ok(activity) => match activity.payload {
+                        // Correct variant, convert to work-specific type
+                        ActivityPayload::RebaseRequest(req) => {
+                            Poll::Ready(Some(Ok(AckRebaseRequest {
+                                id: activity.id,
+                                payload: req,
+                                metadata: activity.metadata,
+                                acker: Arc::new(acker),
+                            })))
+                        }
+                        // Unexpected variant, message is invalid
+                        _ => Poll::Ready(Some(Err(LayerDbError::UnexpectedActivityVariant(
+                            ActivityPayloadDiscriminants::RebaseRequest.to_subject(),
+                            ActivityPayloadDiscriminants::from(activity.payload).to_subject(),
+                        )))),
+                    },
+                    // Error deserializing message
+                    Err(err) => Poll::Ready(Some(Err(err.into()))),
+                }
+            }
+            // Upstream errors are propagated downstream
+            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err.into()))),
+            // If the upstream closes, then we do too
+            Poll::Ready(None) => Poll::Ready(None),
+            // Not ready, so...not ready!
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}

lib/si-layer-cache/src/db.rs

@@ -9,7 +9,10 @@ use tokio_util::{sync::CancellationToken, task::TaskTracker};
 use ulid::Ulid;
 
 use crate::{
-    activities::{Activity, ActivityPayloadDiscriminants, ActivityPublisher, ActivityStream},
+    activities::{
+        Activity, ActivityPayloadDiscriminants, ActivityPublisher, ActivityStream,
+        RebaserRequestsWorkQueueStream,
+    },
     error::LayerDbResult,
     layer_cache::LayerCache,
     persister::{PersisterClient, PersisterTask},
@@ -168,6 +171,12 @@ where
         )
         .await
     }
+
+    pub async fn subscribe_rebaser_requests_work_queue(
+        &self,
+    ) -> LayerDbResult<RebaserRequestsWorkQueueStream> {
+        RebaserRequestsWorkQueueStream::create(&self.nats_client).await
+    }
 }
 
 #[must_use = "graceful shutdown must be spawned on runtime"]

lib/si-layer-cache/src/error.rs

@@ -63,6 +63,8 @@ pub enum LayerDbError {
     SledError(#[from] sled::Error),
     #[error("tokio oneshot recv error: {0}")]
     TokioOneShotRecv(#[from] tokio::sync::oneshot::error::RecvError),
+    #[error("unexpected activity variant; expected={0}, actual={1}")]
+    UnexpectedActivityVariant(String, String),
 }
 
 impl LayerDbError {

lib/si-layer-cache/src/nats.rs

@@ -17,6 +17,8 @@ const NATS_EVENT_STREAM_SUBJECTS: &[&str] = &["si.layerdb.events.*.*.*.*"];
 const NATS_ACTIVITIES_STREAM_NAME: &str = "LAYERDB_ACTIVITIES";
 const NATS_ACTIVITIES_STREAM_SUBJECTS: &[&str] = &["si.layerdb.activities.>"];
 
+const NATS_REBASER_REQUESTS_WORK_QUEUE_STREAM_NAME: &str = "REBASER_REQUESTS";
+
 /// Returns a Jetstream Stream and creates it if it doesn't yet exist.
 pub async fn layerdb_events_stream(
     context: &jetstream::Context,
@@ -68,6 +70,34 @@ pub async fn layerdb_activities_stream(
     Ok(stream)
 }
 
+pub async fn rebaser_requests_work_queue_stream(
+    context: &jetstream::Context,
+    prefix: Option<&str>,
+) -> Result<jetstream::stream::Stream, jetstream::context::CreateStreamError> {
+    let requests_subject = subject::for_activity_discriminate(
+        prefix,
+        crate::activities::ActivityPayloadDiscriminants::RebaseRequest,
+    );
+
+    let source = jetstream::stream::Source {
+        name: nats_stream_name(prefix, NATS_ACTIVITIES_STREAM_NAME),
+        filter_subject: Some(requests_subject.to_string()),
+        ..Default::default()
+    };
+
+    let stream = context
+        .get_or_create_stream(jetstream::stream::Config {
+            name: nats_stream_name(prefix, NATS_REBASER_REQUESTS_WORK_QUEUE_STREAM_NAME),
+            description: Some("Rebaser requests work queue".to_owned()),
+            retention: jetstream::stream::RetentionPolicy::WorkQueue,
+            sources: Some(vec![source]),
+            ..Default::default()
+        })
+        .await?;
+
+    Ok(stream)
+}
+
 fn nats_stream_name(prefix: Option<&str>, suffix: impl AsRef<str>) -> String {
     let suffix = suffix.as_ref();