Pass Headers in call/send

Author: igalshilman

python/restate/server_context.py

@@ -219,7 +219,7 @@ async def create_poll_or_cancel_coroutine(self, handle) -> bytes | None:
                 return self.must_take_notification(handle)
 
             # Nothing ready yet, let's try to make some progress
-            do_progress_response = self.vm.do_progress([handle, CANCEL_HANDLE])
+            do_progress_response = self.vm.do_progress([handle])
             if isinstance(do_progress_response, PyDoProgressAnyCompleted):
                 # One of the handles completed, we can continue
                 continue
@@ -339,14 +339,16 @@ def do_call(self,
                 key: Optional[str] = None,
                 send_delay: Optional[timedelta] = None,
                 send: bool = False,
-                idempotency_key: str | None = None) -> Awaitable[O] | SendHandle:
+                idempotency_key: str | None = None,
+                headers: typing.List[typing.Tuple[str, str]] | None = None
+                ) -> Awaitable[O] | SendHandle:
         """Make an RPC call to the given handler"""
         target_handler = handler_from_callable(tpe)
         service=target_handler.service_tag.name
         handler=target_handler.name
         input_serde = target_handler.handler_io.input_serde
         output_serde = target_handler.handler_io.output_serde
-        return self.do_raw_call(service, handler, parameter, input_serde, output_serde, key, send_delay, send, idempotency_key)
+        return self.do_raw_call(service, handler, parameter, input_serde, output_serde, key, send_delay, send, idempotency_key, headers)
 
 
     def do_raw_call(self,
@@ -358,23 +360,25 @@ def do_raw_call(self,
                  key: Optional[str] = None,
                  send_delay: Optional[timedelta] = None,
                  send: bool = False,
-                 idempotency_key: str | None = None
+                 idempotency_key: str | None = None,
+                 headers: typing.List[typing.Tuple[str, str]] | None = None
                  ) -> Awaitable[O] | SendHandle:
         """Make an RPC call to the given handler"""
         parameter = input_serde.serialize(input_param)
         if send_delay:
             ms = int(send_delay.total_seconds() * 1000)
-            send_handle = self.vm.sys_send(service, handler, parameter, key, delay=ms, idempotency_key=idempotency_key)
+            send_handle = self.vm.sys_send(service, handler, parameter, key, delay=ms, idempotency_key=idempotency_key, headers=headers)
             return ServerSendHandle(self, send_handle)
         if send:
-            send_handle = self.vm.sys_send(service, handler, parameter, key, idempotency_key=idempotency_key)
+            send_handle = self.vm.sys_send(service, handler, parameter, key, idempotency_key=idempotency_key, headers=headers)
             return ServerSendHandle(self, send_handle)
 
         handle = self.vm.sys_call(service=service,
                                   handler=handler,
                                   parameter=parameter,
                                   key=key,
-                                  idempotency_key=idempotency_key)
+                                  idempotency_key=idempotency_key,
+                                  headers=headers)
 
         async def await_point(s: ServerInvocationContext, h, o: Serde[O]):
             """Wait for this handle to be resolved, and deserialize the response."""
@@ -386,52 +390,75 @@ async def await_point(s: ServerInvocationContext, h, o: Serde[O]):
     def service_call(self,
                      tpe: Callable[[Any, I], Awaitable[O]],
                      arg: I,
-                     idempotency_key: str | None = None
+                     idempotency_key: str | None = None,
+                     headers: typing.List[typing.Tuple[str, str]] | None = None
                      ) -> Awaitable[O]:
-        coro = self.do_call(tpe, arg, idempotency_key=idempotency_key)
+        coro = self.do_call(tpe, arg, idempotency_key=idempotency_key, headers=headers)
         assert not isinstance(coro, SendHandle)
         return coro
 
-    def service_send(self, tpe: Callable[[Any, I], Awaitable[O]], arg: I, send_delay: timedelta | None = None, idempotency_key: str | None = None) -> SendHandle:
-        send = self.do_call(tpe=tpe, parameter=arg, send_delay=send_delay, send=True, idempotency_key=idempotency_key)
+    def service_send(self, tpe: Callable[[Any, I], Awaitable[O]], arg: I, send_delay: timedelta | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> SendHandle:
+        send = self.do_call(tpe=tpe, parameter=arg, send_delay=send_delay, send=True, idempotency_key=idempotency_key, headers=headers)
         assert isinstance(send, SendHandle)
         return send
 
     def object_call(self,
                     tpe: Callable[[Any, I],Awaitable[O]],
                     key: str,
                     arg: I,
-                    idempotency_key: str | None = None
+                    idempotency_key: str | None = None,
+                    headers: typing.List[typing.Tuple[str, str]] | None = None
                     ) -> Awaitable[O]:
-        coro = self.do_call(tpe, arg, key, idempotency_key=idempotency_key)
+        coro = self.do_call(tpe, arg, key, idempotency_key=idempotency_key, headers=headers)
         assert not isinstance(coro, SendHandle)
         return coro
 
-    def object_send(self, tpe: Callable[[Any, I], Awaitable[O]], key: str, arg: I, send_delay: timedelta | None = None, idempotency_key: str | None = None) -> SendHandle:
-        send = self.do_call(tpe=tpe, key=key, parameter=arg, send_delay=send_delay, send=True, idempotency_key=idempotency_key)
+    def object_send(self, tpe: Callable[[Any, I], Awaitable[O]], key: str, arg: I, send_delay: timedelta | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> SendHandle:
+        send = self.do_call(tpe=tpe, key=key, parameter=arg, send_delay=send_delay, send=True, idempotency_key=idempotency_key, headers=headers)
         assert isinstance(send, SendHandle)
         return send
 
     def workflow_call(self,
                         tpe: Callable[[Any, I], Awaitable[O]],
                         key: str,
                         arg: I,
-                        idempotency_key: str | None = None
+                        idempotency_key: str | None = None,
+                        headers: typing.List[typing.Tuple[str, str]] | None = None
                         ) -> Awaitable[O]:
-        return self.object_call(tpe, key, arg, idempotency_key=idempotency_key)
+        return self.object_call(tpe, key, arg, idempotency_key=idempotency_key, headers=headers)
 
-    def workflow_send(self, tpe: Callable[[Any, I], Awaitable[O]], key: str, arg: I, send_delay: timedelta | None = None, idempotency_key: str | None = None) -> SendHandle:
-        send = self.object_send(tpe, key, arg, send_delay, idempotency_key=idempotency_key)
+    def workflow_send(self, tpe: Callable[[Any, I], Awaitable[O]], key: str, arg: I, send_delay: timedelta | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> SendHandle:
+        send = self.object_send(tpe, key, arg, send_delay, idempotency_key=idempotency_key, headers=headers)
         assert isinstance(send, SendHandle)
         return send
 
-    def generic_call(self, service: str, handler: str, arg: bytes, key: str | None = None, idempotency_key: str | None = None) -> Awaitable[bytes]:
+    def generic_call(self, service: str, handler: str, arg: bytes, key: str | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> Awaitable[bytes]:
         serde = BytesSerde()
-        return self.do_raw_call(service, handler, arg, serde, serde, key, idempotency_key) # type: ignore
-
-    def generic_send(self, service: str, handler: str, arg: bytes, key: str | None = None, send_delay: timedelta | None = None, idempotency_key: str | None = None) -> SendHandle:
+        call_handle = self.do_raw_call(service=service,
+                                handler=handler,
+                                input_param=arg,
+                                input_serde=serde,
+                                output_serde=serde,
+                                key=key,
+                                idempotency_key=idempotency_key,
+                                headers=headers)
+        assert not isinstance(call_handle, SendHandle)
+        return call_handle
+
+    def generic_send(self, service: str, handler: str, arg: bytes, key: str | None = None, send_delay: timedelta | None = None, idempotency_key: str | None = None, headers: typing.List[typing.Tuple[str, str]] | None = None) -> SendHandle:
         serde = BytesSerde()
-        return self.do_raw_call(service, handler, arg, serde, serde , key, send_delay, True, idempotency_key) # type: ignore
+        send_handle =  self.do_raw_call(service=service,
+                                handler=handler,
+                                input_param=arg,
+                                input_serde=serde,
+                                output_serde=serde,
+                                key=key,
+                                send_delay=send_delay,
+                                send=True,
+                                idempotency_key=idempotency_key,
+                                headers=headers)
+        assert isinstance(send_handle, SendHandle)
+        return send_handle
 
     def awakeable(self,
                   serde: typing.Optional[Serde[I]] = JsonSerde()) -> typing.Tuple[str, Awaitable[Any]]:

python/restate/vm.py

@@ -16,7 +16,7 @@
 
 from dataclasses import dataclass
 import typing
-from restate._internal import PyVM, PyFailure, PySuspended, PyVoid, PyStateKeys, PyExponentialRetryConfig, PyDoProgressAnyCompleted, PyDoProgressReadFromInput, PyDoProgressExecuteRun, PyDoProgressCancelSignalReceived, CANCEL_NOTIFICATION_HANDLE  # pylint: disable=import-error,no-name-in-module,line-too-long
+from restate._internal import PyVM, PyHeader, PyFailure, PySuspended, PyVoid, PyStateKeys, PyExponentialRetryConfig, PyDoProgressAnyCompleted, PyDoProgressReadFromInput, PyDoProgressExecuteRun, PyDoProgressCancelSignalReceived, CANCEL_NOTIFICATION_HANDLE  # pylint: disable=import-error,no-name-in-module,line-too-long
 
 @dataclass
 class Invocation:
@@ -245,10 +245,13 @@ def sys_call(self,
                  handler: str,
                  parameter: bytes,
                  key: typing.Optional[str] = None,
-                 idempotency_key: typing.Optional[str] = None
+                 idempotency_key: typing.Optional[str] = None,
+                 headers: typing.Optional[typing.List[typing.Tuple[str, str]]] = None
                  ):
         """Call a service"""
-        return self.vm.sys_call(service, handler, parameter, key, idempotency_key)
+        if headers:
+            headers = [PyHeader(key=h[0], value=h[1]) for h in headers]
+        return self.vm.sys_call(service, handler, parameter, key, idempotency_key, headers)
 
     # pylint: disable=too-many-arguments
     def sys_send(self,
@@ -257,13 +260,16 @@ def sys_send(self,
                  parameter: bytes,
                  key: typing.Optional[str] = None,
                  delay: typing.Optional[int] = None,
-                 idempotency_key: typing.Optional[str] = None
+                 idempotency_key: typing.Optional[str] = None,
+                 headers: typing.Optional[typing.List[typing.Tuple[str, str]]] = None
                  ) -> int:
         """
         send an invocation to a service, and return the handle
         to the promise that will resolve with the invocation id
         """
-        return self.vm.sys_send(service, handler, parameter, key, delay, idempotency_key)
+        if headers:
+            headers = [PyHeader(key=h[0], value=h[1]) for h in headers]
+        return self.vm.sys_send(service, handler, parameter, key, delay, idempotency_key, headers)
 
     def sys_run(self, name: str) -> int:
         """

src/lib.rs

@@ -22,6 +22,14 @@ struct PyHeader {
     value: String,
 }
 
+#[pymethods]
+impl PyHeader {
+    #[new]
+    fn new(key: String, value: String) -> PyHeader {
+        Self { key, value }
+    }
+}
+
 impl From<Header> for PyHeader {
     fn from(h: Header) -> Self {
         PyHeader {
@@ -31,6 +39,15 @@ impl From<Header> for PyHeader {
     }
 }
 
+impl From<PyHeader> for Header {
+    fn from(h: PyHeader) -> Self {
+        Header {
+            key: h.key.into(),
+            value: h.value.into(),
+        }
+    }
+}
+
 #[pyclass]
 struct PyResponseHead {
     #[pyo3(get, set)]
@@ -433,14 +450,15 @@ impl PyVM {
             .map_err(Into::into)
     }
 
-    #[pyo3(signature = (service, handler, buffer, key=None, idempotency_key=None))]
+    #[pyo3(signature = (service, handler, buffer, key=None, idempotency_key=None, headers=None))]
     fn sys_call(
         mut self_: PyRefMut<'_, Self>,
         service: String,
         handler: String,
         buffer: &Bound<'_, PyBytes>,
         key: Option<String>,
         idempotency_key: Option<String>,
+        headers: Option<Vec<PyHeader>>,
     ) -> Result<PyCallHandle, PyVMError> {
         self_
             .vm
@@ -450,15 +468,19 @@ impl PyVM {
                     handler,
                     key,
                     idempotency_key,
-                    headers: vec![],
+                    headers: headers
+                        .unwrap_or_default()
+                        .into_iter()
+                        .map(Into::into)
+                        .collect(),
                 },
                 buffer.as_bytes().to_vec().into(),
             )
             .map(Into::into)
             .map_err(Into::into)
     }
 
-    #[pyo3(signature = (service, handler, buffer, key=None, delay=None, idempotency_key=None))]
+    #[pyo3(signature = (service, handler, buffer, key=None, delay=None, idempotency_key=None, headers=None))]
     fn sys_send(
         mut self_: PyRefMut<'_, Self>,
         service: String,
@@ -467,6 +489,7 @@ impl PyVM {
         key: Option<String>,
         delay: Option<u64>,
         idempotency_key: Option<String>, 
+        headers: Option<Vec<PyHeader>>,
     ) -> Result<PyNotificationHandle, PyVMError> {
         self_
             .vm
@@ -476,7 +499,11 @@ impl PyVM {
                     handler,
                     key,
                     idempotency_key,
-                    headers: vec![],
+                    headers: headers
+                        .unwrap_or_default()
+                        .into_iter()
+                        .map(Into::into)
+                        .collect(),
                 },
                 buffer.as_bytes().to_vec().into(),
                 delay.map(|millis| {