Support headers in service to service call/send

restatedev · Dec 6, 2024 · b6285c6 · b6285c6
1 parent b6e7963
commit b6285c6
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 9 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -31,6 +31,24 @@ pub struct Header {
     pub value: Cow<'static, str>,
 }
 
+impl From<service_protocol::messages::Header> for Header {
+    fn from(header: service_protocol::messages::Header) -> Self {
+        Self {
+            key: Cow::Owned(header.key),
+            value: Cow::Owned(header.value),
+        }
+    }
+}
+
+impl From<Header> for service_protocol::messages::Header {
+    fn from(header: Header) -> Self {
+        Self {
+            key: header.key.into(),
+            value: header.value.into(),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct ResponseHead {
     pub status_code: u16,
@@ -123,6 +141,7 @@ pub struct Target {
     pub handler: String,
     pub key: Option<String>,
     pub idempotency_key: Option<String>,
+    pub headers: Vec<Header>,
 }
 
 #[derive(Debug, Hash, Clone, Copy, Eq, PartialEq)]

diff --git a/src/tests/async_result.rs b/src/tests/async_result.rs
@@ -11,6 +11,7 @@ fn greeter_target() -> Target {
         handler: "greeter".to_string(),
         key: None,
         idempotency_key: None,
+        headers: Vec::new(),
     }
 }
 

diff --git a/src/tests/calls.rs b/src/tests/calls.rs
@@ -26,6 +26,7 @@ fn call_then_get_invocation_id_then_cancel_invocation() {
                         handler: "MyHandler".to_string(),
                         key: None,
                         idempotency_key: None,
+                        headers: Vec::new(),
                     },
                     Bytes::new(),
                 )
@@ -113,6 +114,7 @@ fn send_then_get_invocation_id_then_cancel_invocation() {
                         handler: "MyHandler".to_string(),
                         key: None,
                         idempotency_key: None,
+                        headers: Vec::new(),
                     },
                     Bytes::new(),
                     None,

diff --git a/src/tests/failures.rs b/src/tests/failures.rs
@@ -76,6 +76,7 @@ fn one_way_call_entry_mismatch() {
                     handler: "greet".to_owned(),
                     key: Some("my-key".to_owned()),
                     idempotency_key: None,
+                    headers: Vec::new(),
                 },
                 Bytes::from_static(b"456"),
                 None,

diff --git a/src/vm/mod.rs b/src/vm/mod.rs
@@ -456,6 +456,11 @@ impl super::VM for CoreVM {
                 handler_name: target.handler,
                 key: target.key.unwrap_or_default(),
                 idempotency_key: target.idempotency_key,
+                headers: target
+                    .headers
+                    .into_iter()
+                    .map(crate::service_protocol::messages::Header::from)
+                    .collect(),
                 parameter: input,
                 ..Default::default()
             },
@@ -491,6 +496,11 @@ impl super::VM for CoreVM {
                 handler_name: target.handler,
                 key: target.key.unwrap_or_default(),
                 idempotency_key: target.idempotency_key,
+                headers: target
+                    .headers
+                    .into_iter()
+                    .map(crate::service_protocol::messages::Header::from)
+                    .collect(),
                 parameter: input,
                 invoke_time: delay
                     .map(|d| {

diff --git a/src/vm/transitions/journal.rs b/src/vm/transitions/journal.rs
@@ -1,5 +1,4 @@
 use crate::retries::NextRetry;
-use crate::service_protocol::messages;
 use crate::service_protocol::messages::{
     run_entry_message, CompletableEntryMessage, EntryMessage, EntryMessageHeaderEq,
     InputEntryMessage, RestateMessage, RunEntryMessage, WriteableRestateMessage,
@@ -121,14 +120,7 @@ impl TransitionAndReturn<Context, SysInput> for State {
                 invocation_id: start_info.debug_id.clone(),
                 random_seed: compute_random_seed(&start_info.id),
                 key: start_info.key.clone(),
-                headers: msg
-                    .headers
-                    .into_iter()
-                    .map(|messages::Header { key, value }| Header {
-                        key: key.into(),
-                        value: value.into(),
-                    })
-                    .collect(),
+                headers: msg.headers.into_iter().map(Header::from).collect(),
                 input: msg.value,
             },
         ))