Few changes to the replay aware filter feature:

* We now emit the span all the times for the endpoint handle method. This seems a reasonable default, and doesn't require the tracing-subscriber dependency (we already pull in tracing as dependency anyway)
* Renamed the span name and the replaying field filter, to avoid confusion and clash
* Now the replay aware filter will look for the exact span name created by the handler. This should make the matching of the exact span we wanna use for filtering more robust.
* Flip the replaying field not only when taking an awaiting, but also when executing `sys_` functions on the state machine. Those are the ones that can cause state transitions.
* Simplify the tracing example

Author: slinkydeveloper

Cargo.toml

@@ -10,14 +10,13 @@ rust-version = "1.76.0"
 [[example]]
 name = "tracing"
 path = "examples/tracing.rs"
-required-features = ["tracing-subscriber"]
+required-features = ["span-filter"]
 
 [features]
-default = ["http_server", "rand", "uuid"]
+default = ["http_server", "rand", "uuid", "span-filter"]
 hyper = ["dep:hyper", "http-body-util", "restate-sdk-shared-core/http"]
 http_server = ["hyper", "hyper/server", "hyper/http2", "hyper-util", "tokio/net", "tokio/signal", "tokio/macros"]
-tracing-subscriber = ["dep:tracing-subscriber"]
-
+span-filter = ["dep:tracing-subscriber"]
 
 [dependencies]
 bytes = "1.6.1"

examples/tracing.rs

@@ -1,51 +1,28 @@
 use restate_sdk::prelude::*;
-use std::convert::Infallible;
 use std::time::Duration;
 use tracing::info;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer};
 
 #[restate_sdk::service]
 trait Greeter {
-    async fn greet(name: String) -> Result<String, Infallible>;
+    async fn greet(name: String) -> Result<String, HandlerError>;
 }
 
 struct GreeterImpl;
 
 impl Greeter for GreeterImpl {
-    async fn greet(&self, ctx: Context<'_>, name: String) -> Result<String, Infallible> {
-        let timeout = 60; // More than suspension timeout to trigger replay
-        info!("This will be logged on replay");
-        _ = ctx.service_client::<DelayerClient>().delay(1).call().await;
-        info!("This will not be logged on replay");
-        _ = ctx
-            .service_client::<DelayerClient>()
-            .delay(timeout)
-            .call()
-            .await;
-        info!("This will be logged on processing after suspension");
-        Ok(format!("Greetings {name} after {timeout} seconds"))
-    }
-}
-
-#[restate_sdk::service]
-trait Delayer {
-    async fn delay(seconds: u64) -> Result<String, Infallible>;
-}
-
-struct DelayerImpl;
-
-impl Delayer for DelayerImpl {
-    async fn delay(&self, ctx: Context<'_>, seconds: u64) -> Result<String, Infallible> {
-        _ = ctx.sleep(Duration::from_secs(seconds)).await;
-        info!("Delayed for {seconds} seconds");
-        Ok(format!("Delayed {seconds}"))
+    async fn greet(&self, ctx: Context<'_>, name: String) -> Result<String, HandlerError> {
+        info!("Before sleep");
+        ctx.sleep(Duration::from_secs(61)).await?; // More than suspension timeout to trigger replay
+        info!("After sleep");
+        Ok(format!("Greetings {name}"))
     }
 }
 
 #[tokio::main]
 async fn main() {
     let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
-        .unwrap_or_else(|_| "restate_sdk=info".into());
+        .unwrap_or_else(|_| "info,restate_sdk=debug".into());
     let replay_filter = restate_sdk::filter::ReplayAwareFilter;
     tracing_subscriber::registry()
         .with(
@@ -54,12 +31,7 @@ async fn main() {
                 .with_filter(replay_filter),
         )
         .init();
-    HttpServer::new(
-        Endpoint::builder()
-            .bind(GreeterImpl.serve())
-            .bind(DelayerImpl.serve())
-            .build(),
-    )
-    .listen_and_serve("0.0.0.0:9080".parse().unwrap())
-    .await;
+    HttpServer::new(Endpoint::builder().bind(GreeterImpl.serve()).build())
+        .listen_and_serve("0.0.0.0:9080".parse().unwrap())
+        .await;
 }

src/endpoint/context.rs

@@ -30,10 +30,10 @@ pub struct ContextInternalInner {
     pub(crate) read: InputReceiver,
     pub(crate) write: OutputSender,
     pub(super) handler_state: HandlerStateNotifier,
-    // Flag to indicate whether span replay attribute should be set
-    // When replaying this is set on the sys call
-    // When not replaying this is reset on the sys call that transitioned the state
-    pub(super) tracing_replaying_flag: bool,
+
+    /// We remember here the state of the span replaying field state, because setting it might be expensive (it's guarded behind locks and other stuff).
+    /// For details, see [ContextInternalInner::maybe_flip_span_replaying_field]
+    pub(super) span_replaying_field_state: bool,
 }
 
 impl ContextInternalInner {
@@ -48,11 +48,12 @@ impl ContextInternalInner {
             read,
             write,
             handler_state,
-            tracing_replaying_flag: true,
+            span_replaying_field_state: false,
         }
     }
 
     pub(super) fn fail(&mut self, e: Error) {
+        self.maybe_flip_span_replaying_field();
         self.vm.notify_error(
             CoreError::new(500u16, e.0.to_string())
                 .with_stacktrace(Cow::Owned(format!("{:#}", e.0))),
@@ -61,19 +62,13 @@ impl ContextInternalInner {
         self.handler_state.mark_error(e);
     }
 
-    pub(super) fn set_tracing_replaying_flag(&mut self) {
-        if !self.vm.is_processing() {
-            // Replay record is not yet set in the span
-            if self.tracing_replaying_flag {
-                tracing::Span::current().record("replaying", true);
-                self.tracing_replaying_flag = false;
-            }
-        } else {
-            // Replay record is not yet reset in the span
-            if !self.tracing_replaying_flag {
-                tracing::Span::current().record("replaying", false);
-                self.tracing_replaying_flag = true;
-            }
+    pub(super) fn maybe_flip_span_replaying_field(&mut self) {
+        if !self.span_replaying_field_state && self.vm.is_replaying() {
+            tracing::Span::current().record("restate.sdk.is_replaying", true);
+            self.span_replaying_field_state = true;
+        } else if self.span_replaying_field_state && !self.vm.is_replaying() {
+            tracing::Span::current().record("restate.sdk.is_replaying", false);
+            self.span_replaying_field_state = false;
         }
     }
 }
@@ -211,6 +206,7 @@ impl ContextInternal {
                         },
                     ))
                 });
+        inner_lock.maybe_flip_span_replaying_field();
 
         match input_result {
             Ok(Ok(i)) => {
@@ -244,6 +240,7 @@ impl ContextInternal {
     ) -> impl Future<Output = Result<Option<T>, TerminalError>> + Send {
         let mut inner_lock = must_lock!(self.inner);
         let handle = unwrap_or_trap!(inner_lock, inner_lock.vm.sys_state_get(key.to_owned()));
+        inner_lock.maybe_flip_span_replaying_field();
 
         let poll_future = get_async_result(Arc::clone(&self.inner), handle).map(|res| match res {
             Ok(Value::Void) => Ok(Ok(None)),
@@ -267,6 +264,7 @@ impl ContextInternal {
     pub fn get_keys(&self) -> impl Future<Output = Result<Vec<String>, TerminalError>> + Send {
         let mut inner_lock = must_lock!(self.inner);
         let handle = unwrap_or_trap!(inner_lock, inner_lock.vm.sys_state_get_keys());
+        inner_lock.maybe_flip_span_replaying_field();
 
         let poll_future = get_async_result(Arc::clone(&self.inner), handle).map(|res| match res {
             Ok(Value::Failure(f)) => Ok(Err(f.into())),
@@ -287,6 +285,7 @@ impl ContextInternal {
         match t.serialize() {
             Ok(b) => {
                 let _ = inner_lock.vm.sys_state_set(key.to_owned(), b);
+                inner_lock.maybe_flip_span_replaying_field();
             }
             Err(e) => {
                 inner_lock.fail(Error::serialization("set_state", e));
@@ -295,11 +294,15 @@ impl ContextInternal {
     }
 
     pub fn clear(&self, key: &str) {
-        let _ = must_lock!(self.inner).vm.sys_state_clear(key.to_string());
+        let mut inner_lock = must_lock!(self.inner);
+        let _ = inner_lock.vm.sys_state_clear(key.to_string());
+        inner_lock.maybe_flip_span_replaying_field();
     }
 
     pub fn clear_all(&self) {
-        let _ = must_lock!(self.inner).vm.sys_state_clear_all();
+        let mut inner_lock = must_lock!(self.inner);
+        let _ = inner_lock.vm.sys_state_clear_all();
+        inner_lock.maybe_flip_span_replaying_field();
     }
 
     pub fn sleep(
@@ -314,6 +317,7 @@ impl ContextInternal {
             inner_lock,
             inner_lock.vm.sys_sleep(now + sleep_duration, Some(now))
         );
+        inner_lock.maybe_flip_span_replaying_field();
 
         let poll_future = get_async_result(Arc::clone(&self.inner), handle).map(|res| match res {
             Ok(Value::Void) => Ok(Ok(())),
@@ -349,6 +353,7 @@ impl ContextInternal {
         );
 
         let call_handle = unwrap_or_trap!(inner_lock, inner_lock.vm.sys_call(target, input));
+        inner_lock.maybe_flip_span_replaying_field();
         drop(inner_lock);
 
         // Let's prepare the two futures here
@@ -432,6 +437,7 @@ impl ContextInternal {
                 return Either::Right(TrapFuture::<()>::default());
             }
         };
+        inner_lock.maybe_flip_span_replaying_field();
         drop(inner_lock);
 
         let invocation_id_fut = InterceptErrorFuture::new(
@@ -473,6 +479,7 @@ impl ContextInternal {
     ) {
         let mut inner_lock = must_lock!(self.inner);
         let maybe_awakeable_id_and_handle = inner_lock.vm.sys_awakeable();
+        inner_lock.maybe_flip_span_replaying_field();
 
         let (awakeable_id, handle) = match maybe_awakeable_id_and_handle {
             Ok((s, handle)) => (s, handle),
@@ -533,6 +540,7 @@ impl ContextInternal {
     ) -> impl Future<Output = Result<T, TerminalError>> + Send {
         let mut inner_lock = must_lock!(self.inner);
         let handle = unwrap_or_trap!(inner_lock, inner_lock.vm.sys_get_promise(name.to_owned()));
+        inner_lock.maybe_flip_span_replaying_field();
         drop(inner_lock);
 
         let poll_future = get_async_result(Arc::clone(&self.inner), handle).map(|res| match res {
@@ -558,6 +566,7 @@ impl ContextInternal {
     ) -> impl Future<Output = Result<Option<T>, TerminalError>> + Send {
         let mut inner_lock = must_lock!(self.inner);
         let handle = unwrap_or_trap!(inner_lock, inner_lock.vm.sys_peek_promise(name.to_owned()));
+        inner_lock.maybe_flip_span_replaying_field();
         drop(inner_lock);
 
         let poll_future = get_async_result(Arc::clone(&self.inner), handle).map(|res| match res {
@@ -646,6 +655,7 @@ impl ContextInternal {
         };
 
         let _ = inner_lock.vm.sys_write_output(res_to_write);
+        inner_lock.maybe_flip_span_replaying_field();
     }
 
     pub fn end(&self) {
@@ -880,6 +890,7 @@ impl<InvIdFut: Future<Output = Result<String, TerminalError>> + Send> Invocation
             let inv_id = cloned_invocation_id_fut.await?;
             let mut inner_lock = must_lock!(cloned_ctx);
             let _ = inner_lock.vm.sys_cancel_invocation(inv_id);
+            inner_lock.maybe_flip_span_replaying_field();
             drop(inner_lock);
             Ok(())
         }
@@ -924,6 +935,7 @@ where
             let inv_id = cloned_invocation_id_fut.await?;
             let mut inner_lock = must_lock!(cloned_ctx);
             let _ = inner_lock.vm.sys_cancel_invocation(inv_id);
+            inner_lock.maybe_flip_span_replaying_field();
             drop(inner_lock);
             Ok(())
         }

src/endpoint/futures/async_result_poll.rs

@@ -128,6 +128,9 @@ impl Future for VmAsyncResultPollFuture {
                         }
                     };
 
+                    // DoProgress might cause a flip of the replaying state
+                    inner_lock.maybe_flip_span_replaying_field();
+
                     // At this point let's try to take the notification
                     match inner_lock.vm.take_notification(handle) {
                         Ok(Some(v)) => return Poll::Ready(Ok(v)),

src/endpoint/mod.rs

@@ -18,8 +18,7 @@ use std::future::poll_fn;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
-#[cfg(feature = "tracing-subscriber")]
-use tracing::{field, info_span, Instrument};
+use tracing::{info_span, Instrument};
 
 const DISCOVERY_CONTENT_TYPE: &str = "application/vnd.restate.endpointmanifest.v1+json";
 
@@ -370,27 +369,13 @@ impl BidiStreamRunner {
             .get(&self.svc_name)
             .expect("service must exist at this point");
 
-        #[cfg(feature = "tracing-subscriber")]
-        {
-            let span = info_span!(
-                "handle",
-                "rpc.system" = "restate",
-                "rpc.service" = self.svc_name,
-                "rpc.method" = self.handler_name,
-                "replaying" = field::Empty,
-            );
-            handle(
-                input_rx,
-                output_tx,
-                self.vm,
-                self.svc_name,
-                self.handler_name,
-                svc,
-            )
-            .instrument(span)
-            .await
-        }
-        #[cfg(not(feature = "tracing-subscriber"))]
+        let span = info_span!(
+            "restate_sdk_endpoint_handle",
+            "rpc.system" = "restate",
+            "rpc.service" = self.svc_name,
+            "rpc.method" = self.handler_name,
+            "restate.sdk.is_replaying" = false
+        );
         handle(
             input_rx,
             output_tx,
@@ -399,6 +384,7 @@ impl BidiStreamRunner {
             self.handler_name,
             svc,
         )
+        .instrument(span)
         .await
     }
 }