restate-sdk 0.11.0 → 0.12.0

data/lib/restate/server/context.rb

@@ -0,0 +1,680 @@
+# typed: false
+# frozen_string_literal: true
+
+require 'async'
+require 'async/queue'
+require 'logger'
+
+module Restate
+  module Server
+    # The core execution context for a Restate handler invocation.
+    # Implements the progress loop and all context API methods (state, run, sleep, call, send).
+    #
+    # Concurrency model:
+    #   - The handler runs inside a Fiber managed by Falcon/Async.
+    #   - `run` blocks spawn child Async tasks.
+    #   - When the progress loop needs input, it dequeues from @input_queue, yielding the Fiber.
+    #   - The HTTP input reader (a separate Async task) feeds chunks into @input_queue.
+    #   - Output chunks are written directly to the streaming response body.
+    class Context
+      include WorkflowContext
+      include WorkflowSharedContext
+
+      LOGGER = Logger.new($stdout, progname: 'Restate::Server::Context')
+
+      attr_reader :vm, :invocation
+
+      def initialize(vm:, handler:, invocation:, send_output:, input_queue:, middleware: [],
+                     outbound_middleware: [])
+        @vm = vm
+        @handler = handler
+        @invocation = invocation
+        @send_output = send_output
+        @input_queue = input_queue
+        @run_coros_to_execute = {}
+        @attempt_finished_event = AttemptFinishedEvent.new
+        @middleware = middleware
+        @outbound_middleware = outbound_middleware
+      end
+
+      # ── Main entry point ──
+
+      # Runs the handler to completion, writing the output (or failure) to the journal.
+      def enter
+        Thread.current[:restate_context] = self
+        Thread.current[:restate_service_kind] = @handler.service_tag.kind
+        Thread.current[:restate_handler_kind] = @handler.kind
+        in_buffer = @invocation.input_buffer
+        out_buffer = Restate.invoke_handler(handler: @handler, ctx: self, in_buffer: in_buffer,
+                                            middleware: @middleware)
+        @vm.sys_write_output_success(out_buffer)
+        @vm.sys_end
+      rescue TerminalError => e
+        failure = Failure.new(code: e.status_code, message: e.message)
+        @vm.sys_write_output_failure(failure)
+        @vm.sys_end
+      rescue SuspendedError, InternalError
+        # These are expected internal control flow exceptions; do nothing.
+      rescue DisconnectedError
+        raise
+      rescue StandardError => e
+        # Walk the cause chain for TerminalError or internal exceptions
+        cause = e
+        handled = false
+        while cause
+          if cause.is_a?(TerminalError)
+            f = Failure.new(code: cause.status_code, message: cause.message)
+            @vm.sys_write_output_failure(f)
+            @vm.sys_end
+            handled = true
+            break
+          elsif cause.is_a?(SuspendedError) || cause.is_a?(InternalError)
+            handled = true
+            break
+          end
+          cause = cause.cause
+        end
+        unless handled
+          @vm.notify_error(e.inspect, e.backtrace&.join("\n"))
+          raise
+        end
+      ensure
+        @run_coros_to_execute.clear
+        Thread.current[:restate_context] = nil
+        Thread.current[:restate_service_kind] = nil
+        Thread.current[:restate_handler_kind] = nil
+      end
+
+      # Called by the server when the attempt ends (handler completed, disconnected,
+      # or transient error). Signals the attempt_finished_event so that user code
+      # and background pool jobs can clean up.
+      def on_attempt_finished
+        @attempt_finished_event.set!
+      end
+
+      # ── State operations ──
+
+      # Durably retrieves a state entry by name. Returns nil if unset.
+      def get(name, serde: JsonSerde)
+        get_async(name, serde: serde).await
+      end
+
+      # Returns a DurableFuture for a state entry. Resolves to nil if unset.
+      def get_async(name, serde: JsonSerde)
+        handle = @vm.sys_get_state(name)
+        DurableFuture.new(self, handle, serde: serde)
+      end
+
+      # Durably sets a state entry. The value is serialized via +serde+.
+      def set(name, value, serde: JsonSerde)
+        @vm.sys_set_state(name, serde.serialize(value))
+      end
+
+      # Durably removes a single state entry by name.
+      def clear(name)
+        @vm.sys_clear_state(name)
+      end
+
+      # Durably removes all state entries for this virtual object or workflow.
+      def clear_all
+        @vm.sys_clear_all_state
+      end
+
+      # Returns the list of all state entry names for this virtual object or workflow.
+      def state_keys
+        state_keys_async.await
+      end
+
+      # Returns a DurableFuture for the list of all state entry names.
+      def state_keys_async
+        handle = @vm.sys_get_state_keys
+        DurableFuture.new(self, handle)
+      end
+
+      # ── Sleep ──
+
+      # Returns a durable future that completes after the given duration.
+      # The timer survives handler restarts.
+      def sleep(seconds)
+        millis = (seconds * 1000).to_i
+        handle = @vm.sys_sleep(millis)
+        DurableFuture.new(self, handle)
+      end
+
+      # Block until a previously created handle completes. Returns the value.
+      def resolve_handle(handle)
+        poll_and_take(handle)
+      end
+
+      # Wait until any of the given handles completes. Does not take notifications.
+      def wait_any_handle(handles)
+        poll_or_cancel(handles) unless handles.any? { |h| @vm.is_completed(h) }
+      end
+
+      # Check if a handle is completed (non-blocking).
+      def completed?(handle)
+        @vm.is_completed(handle)
+      end
+
+      # Take a completed handle's notification, returning the value.
+      # Raises TerminalError if the handle resolved to a failure.
+      def take_completed(handle)
+        must_take_notification(handle)
+      end
+
+      # Wait until any of the given futures completes. Returns [completed, remaining].
+      def wait_any(*futures)
+        handles = futures.map(&:handle)
+        wait_any_handle(handles)
+        futures.partition(&:completed?)
+      end
+
+      # ── Durable run (side effect) ──
+
+      # Executes a durable side effect. The block runs at most once; its result is
+      # journaled and replayed on retries. Returns a DurableFuture for the result.
+      #
+      # Pass +background: true+ to run the block in a real OS Thread, keeping the
+      # fiber event loop responsive for other concurrent handlers. Use this for
+      # CPU-intensive work.
+      def run(name, serde: JsonSerde, retry_policy: nil, background: false, &action)
+        handle = @vm.sys_run(name)
+
+        @run_coros_to_execute[handle] =
+          if background
+            -> { execute_run_threaded(handle, action, serde, retry_policy) }
+          else
+            -> { execute_run(handle, action, serde, retry_policy) }
+          end
+
+        DurableFuture.new(self, handle, serde: serde)
+      end
+
+      # Convenience shortcut for +run(...).await+ — executes the durable side effect
+      # and returns the result directly.
+      #
+      # Accepts all the same options as +run+, including +background: true+.
+      def run_sync(name, serde: JsonSerde, retry_policy: nil, background: false, &action)
+        run(name, serde: serde, retry_policy: retry_policy, background: background, &action).await
+      end
+
+      # ── Service calls ──
+
+      # Durably calls a handler on a Restate service and returns a future for its result.
+      def service_call(service, handler, arg, key: nil, idempotency_key: nil, headers: nil,
+                       input_serde: NOT_SET, output_serde: NOT_SET)
+        svc_name, handler_name, handler_meta = resolve_call_target(service, handler)
+        in_serde = resolve_serde(input_serde, handler_meta, :input_serde)
+        out_serde = resolve_serde(output_serde, handler_meta, :output_serde)
+        parameter = in_serde.serialize(arg)
+        with_outbound_middleware(svc_name, handler_name, headers, handler_meta: handler_meta) do |hdrs|
+          call_handle = @vm.sys_call(
+            service: svc_name, handler: handler_name, parameter: parameter,
+            key: key, idempotency_key: idempotency_key, headers: hdrs
+          )
+          DurableCallFuture.new(self, call_handle.result_handle, call_handle.invocation_id_handle,
+                                output_serde: out_serde)
+        end
+      end
+
+      # Sends a one-way invocation to a Restate service handler (fire-and-forget).
+      def service_send(service, handler, arg, key: nil, delay: nil, idempotency_key: nil, headers: nil,
+                       input_serde: NOT_SET)
+        svc_name, handler_name, handler_meta = resolve_call_target(service, handler)
+        in_serde = resolve_serde(input_serde, handler_meta, :input_serde)
+        parameter = in_serde.serialize(arg)
+        delay_ms = delay ? (delay * 1000).to_i : nil
+        with_outbound_middleware(svc_name, handler_name, headers, handler_meta: handler_meta) do |hdrs|
+          invocation_id_handle = @vm.sys_send(
+            service: svc_name, handler: handler_name, parameter: parameter,
+            key: key, delay: delay_ms, idempotency_key: idempotency_key, headers: hdrs
+          )
+          SendHandle.new(self, invocation_id_handle)
+        end
+      end
+
+      # Durably calls a handler on a Restate virtual object, keyed by +key+.
+      def object_call(service, handler, key, arg, idempotency_key: nil, headers: nil,
+                      input_serde: NOT_SET, output_serde: NOT_SET)
+        svc_name, handler_name, handler_meta = resolve_call_target(service, handler)
+        in_serde = resolve_serde(input_serde, handler_meta, :input_serde)
+        out_serde = resolve_serde(output_serde, handler_meta, :output_serde)
+        parameter = in_serde.serialize(arg)
+        with_outbound_middleware(svc_name, handler_name, headers, handler_meta: handler_meta) do |hdrs|
+          call_handle = @vm.sys_call(
+            service: svc_name, handler: handler_name, parameter: parameter,
+            key: key, idempotency_key: idempotency_key, headers: hdrs
+          )
+          DurableCallFuture.new(self, call_handle.result_handle, call_handle.invocation_id_handle,
+                                output_serde: out_serde)
+        end
+      end
+
+      # Sends a one-way invocation to a Restate virtual object handler (fire-and-forget).
+      def object_send(service, handler, key, arg, delay: nil, idempotency_key: nil, headers: nil,
+                      input_serde: NOT_SET)
+        svc_name, handler_name, handler_meta = resolve_call_target(service, handler)
+        in_serde = resolve_serde(input_serde, handler_meta, :input_serde)
+        parameter = in_serde.serialize(arg)
+        delay_ms = delay ? (delay * 1000).to_i : nil
+        with_outbound_middleware(svc_name, handler_name, headers, handler_meta: handler_meta) do |hdrs|
+          invocation_id_handle = @vm.sys_send(
+            service: svc_name, handler: handler_name, parameter: parameter,
+            key: key, delay: delay_ms, idempotency_key: idempotency_key, headers: hdrs
+          )
+          SendHandle.new(self, invocation_id_handle)
+        end
+      end
+
+      # Durably calls a handler on a Restate workflow, keyed by +key+.
+      def workflow_call(service, handler, key, arg, idempotency_key: nil, headers: nil,
+                        input_serde: NOT_SET, output_serde: NOT_SET)
+        object_call(service, handler, key, arg, idempotency_key: idempotency_key, headers: headers,
+                    input_serde: input_serde, output_serde: output_serde) # rubocop:disable Layout/HashAlignment
+      end
+
+      # Sends a one-way invocation to a Restate workflow handler (fire-and-forget).
+      def workflow_send(service, handler, key, arg, delay: nil, idempotency_key: nil, headers: nil,
+                        input_serde: NOT_SET)
+        object_send(service, handler, key, arg, delay: delay, idempotency_key: idempotency_key, headers: headers,
+                    input_serde: input_serde) # rubocop:disable Layout/HashAlignment
+      end
+
+      # ── Awakeables ──
+
+      # Creates an awakeable and returns [awakeable_id, DurableFuture].
+      def awakeable(serde: JsonSerde)
+        id, handle = @vm.sys_awakeable
+        [id, DurableFuture.new(self, handle, serde: serde)]
+      end
+
+      # Resolves an awakeable with a success value.
+      def resolve_awakeable(awakeable_id, payload, serde: JsonSerde)
+        @vm.sys_complete_awakeable_success(awakeable_id, serde.serialize(payload))
+      end
+
+      # Rejects an awakeable with a terminal failure.
+      def reject_awakeable(awakeable_id, message, code: 500)
+        failure = Failure.new(code: code, message: message)
+        @vm.sys_complete_awakeable_failure(awakeable_id, failure)
+      end
+
+      # ── Signals ──
+
+      # Wait for a named signal addressed to this invocation. Returns a DurableFuture.
+      def signal(name, serde: JsonSerde)
+        handle = @vm.sys_signal(name)
+        DurableFuture.new(self, handle, serde: serde)
+      end
+
+      # Send a success value to a named signal on another invocation.
+      def resolve_signal(invocation_id, name, payload, serde: JsonSerde)
+        @vm.sys_complete_signal_success(invocation_id, name, serde.serialize(payload))
+      end
+
+      # Send a terminal failure to a named signal on another invocation.
+      def reject_signal(invocation_id, name, message, code: 500)
+        failure = Failure.new(code: code, message: message)
+        @vm.sys_complete_signal_failure(invocation_id, name, failure)
+      end
+
+      # ── Promises (Workflow API) ──
+
+      # Gets a durable promise value, blocking until resolved.
+      def promise(name, serde: JsonSerde)
+        handle = @vm.sys_get_promise(name)
+        poll_and_take(handle) do |raw|
+          raw.nil? ? nil : serde.deserialize(raw)
+        end
+      end
+
+      # Peeks at a durable promise value without blocking. Returns nil if not yet resolved.
+      def peek_promise(name, serde: JsonSerde)
+        handle = @vm.sys_peek_promise(name)
+        poll_and_take(handle) do |raw|
+          raw.nil? ? nil : serde.deserialize(raw)
+        end
+      end
+
+      # Resolves a durable promise with a success value.
+      def resolve_promise(name, payload, serde: JsonSerde)
+        handle = @vm.sys_complete_promise_success(name, serde.serialize(payload))
+        poll_and_take(handle)
+        nil
+      end
+
+      # Rejects a durable promise with a terminal failure.
+      def reject_promise(name, message, code: 500)
+        failure = Failure.new(code: code, message: message)
+        handle = @vm.sys_complete_promise_failure(name, failure)
+        poll_and_take(handle)
+        nil
+      end
+
+      # ── Cancel invocation ──
+
+      # Requests cancellation of another invocation by its id.
+      def cancel_invocation(invocation_id)
+        @vm.sys_cancel_invocation(invocation_id)
+      end
+
+      # ── Generic calls (raw bytes, no serde) ──
+
+      # Durably calls a handler using raw bytes (no serialization). Useful for proxying.
+      def generic_call(service, handler, arg, key: nil, idempotency_key: nil, headers: nil)
+        with_outbound_middleware(service, handler, headers) do |hdrs|
+          call_handle = @vm.sys_call(
+            service: service, handler: handler, parameter: arg,
+            key: key, idempotency_key: idempotency_key, headers: hdrs
+          )
+          DurableCallFuture.new(self, call_handle.result_handle, call_handle.invocation_id_handle,
+                                output_serde: nil)
+        end
+      end
+
+      # Sends a one-way invocation using raw bytes (no serialization). Useful for proxying.
+      def generic_send(service, handler, arg, key: nil, delay: nil, idempotency_key: nil, headers: nil)
+        delay_ms = delay ? (delay * 1000).to_i : nil
+        with_outbound_middleware(service, handler, headers) do |hdrs|
+          invocation_id_handle = @vm.sys_send(
+            service: service, handler: handler, parameter: arg,
+            key: key, delay: delay_ms, idempotency_key: idempotency_key, headers: hdrs
+          )
+          SendHandle.new(self, invocation_id_handle)
+        end
+      end
+
+      # ── Request metadata ──
+
+      # Returns metadata about the current invocation (id, headers, raw body).
+      def request
+        @request ||= Request.new(
+          id: @invocation.invocation_id,
+          headers: @invocation.headers.to_h,
+          body: @invocation.input_buffer,
+          attempt_finished_event: @attempt_finished_event
+        )
+      end
+
+      # Returns the key for this virtual object or workflow invocation.
+      def key
+        @invocation.key
+      end
+
+      private
+
+      # ── Progress loop ──
+
+      # Polls until the given handle(s) complete, then takes the notification.
+      def poll_and_take(handle, &)
+        poll_or_cancel([handle]) unless @vm.is_completed(handle)
+        must_take_notification(handle, &)
+      end
+
+      def poll_or_cancel(handles)
+        loop do
+          flush_output
+          response = @vm.do_progress(handles)
+
+          if response.is_a?(Exception)
+            LOGGER.error("Exception in do_progress: #{response}")
+            raise InternalError
+          end
+
+          case response
+          when Suspended
+            raise SuspendedError
+          when DoProgressAnyCompleted
+            return
+          when DoProgressCancelSignalReceived
+            raise TerminalError.new('cancelled', status_code: 409)
+          when DoProgressExecuteRun
+            fn = @run_coros_to_execute.delete(response.handle)
+            raise "Missing run coroutine for handle #{response.handle}" unless fn
+
+            # Spawn child task for the run action
+            Async do
+              fn.call
+            ensure
+              @input_queue.enqueue(:run_completed)
+            end
+          when DoWaitPendingRun, DoProgressReadFromInput
+            # Wait for input from the HTTP body reader or a run completion signal
+            event = @input_queue.dequeue
+
+            case event
+            when :run_completed
+              next
+            when :eof
+              @vm.notify_input_closed
+            when :disconnected
+              raise DisconnectedError
+            when String
+              @vm.notify_input(event)
+            end
+          end
+        end
+      end
+
+      def must_take_notification(handle, &block)
+        result = @vm.take_notification(handle)
+
+        if result.is_a?(Exception)
+          flush_output
+          LOGGER.error("Exception in take_notification: #{result}")
+          raise InternalError
+        end
+
+        case result
+        when Suspended
+          flush_output
+          raise SuspendedError
+        when NotReady
+          raise "Unexpected NotReady for handle #{handle}"
+        when Failure
+          raise TerminalError.new(result.message, status_code: result.code)
+        else
+          block ? yield(result) : result
+        end
+      end
+
+      def flush_output
+        loop do
+          output = @vm.take_output
+          break if output.nil? || output.empty?
+
+          @send_output.call(output)
+        end
+      end
+
+      # ── Outbound middleware ──
+
+      # Runs outbound middleware chain (Sidekiq client middleware pattern).
+      # Each middleware gets +call(service, handler, headers)+ and must +yield+
+      # to continue the chain. The block at the end performs the actual VM call.
+      #
+      # The optional +handler_meta+ (a Handler struct from resolve_call_target)
+      # is exposed via Thread.current[:restate_outbound_handler_meta] so that
+      # middleware can inspect the target handler's kind without changing the
+      # middleware interface.
+      def with_outbound_middleware(service, handler, headers, handler_meta: nil, &action)
+        return action.call(headers) if @outbound_middleware.empty?
+
+        h = headers || {}
+        previous_meta = Thread.current[:restate_outbound_handler_meta]
+        Thread.current[:restate_outbound_handler_meta] = handler_meta
+        chain = ->(hdrs) { action.call(hdrs) }
+        @outbound_middleware.reverse_each do |mw|
+          prev = chain
+          chain = ->(hdrs) { mw.call(service, handler, hdrs) { prev.call(hdrs) } }
+        end
+        chain.call(h)
+      ensure
+        Thread.current[:restate_outbound_handler_meta] = previous_meta
+      end
+
+      # ── Call target resolution ──
+
+      # Resolves a service+handler pair from class/symbol or string/string.
+      # Returns [service_name, handler_name, handler_metadata_or_nil].
+      def resolve_call_target(service, handler)
+        handler_name = handler.is_a?(Symbol) ? handler.name : handler.to_s
+        if service.is_a?(Class) && service.respond_to?(:service_name)
+          svc_name = service.service_name
+          handler_meta = service.respond_to?(:handlers) ? service.handlers[handler_name] : nil
+          [svc_name, handler_name, handler_meta]
+        else
+          [service.to_s, handler_name, nil]
+        end
+      end
+
+      # Resolves a serde value: if the caller passed NOT_SET, fall back to handler metadata, then JsonSerde.
+      def resolve_serde(caller_serde, handler_meta, field)
+        return caller_serde unless caller_serde.equal?(NOT_SET)
+        return JsonSerde unless handler_meta
+
+        io = handler_meta.handler_io
+        case field
+        when :input_serde then io.input_serde
+        when :output_serde then io.output_serde
+        else JsonSerde
+        end
+      end
+
+      # ── Run execution ──
+
+      def execute_run(handle, action, serde, retry_policy)
+        propose_run_result(handle, action, serde, retry_policy)
+      end
+
+      # Like execute_run, but offloads the action to a real OS Thread.
+      # The fiber yields (via IO.pipe) while the thread runs, keeping the event loop responsive.
+      def execute_run_threaded(handle, action, serde, retry_policy)
+        propose_run_result(handle, -> { offload_to_thread(action) }, serde, retry_policy)
+      end
+
+      # Runs the action and proposes the result (success/failure/transient) to the VM.
+      def propose_run_result(handle, action, serde, retry_policy)
+        start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        begin
+          result = action.call
+          buffer = serde.serialize(result)
+          @vm.propose_run_completion_success(handle, buffer)
+        rescue TerminalError => e
+          failure = Failure.new(code: e.status_code, message: e.message)
+          @vm.propose_run_completion_failure(handle, failure)
+        rescue SuspendedError, InternalError
+          raise
+        rescue StandardError => e
+          elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
+          attempt_duration_ms = (elapsed * 1000).to_i
+          failure = Failure.new(
+            code: 500,
+            message: e.inspect,
+            stacktrace: e.backtrace&.join("\n")
+          )
+          config = RunRetryConfig.new(
+            initial_interval: retry_policy&.initial_interval,
+            max_attempts: retry_policy&.max_attempts,
+            max_duration: retry_policy&.max_duration,
+            max_interval: retry_policy&.max_interval,
+            interval_factor: retry_policy&.interval_factor
+          )
+          @vm.propose_run_completion_transient(
+            handle,
+            failure: failure,
+            attempt_duration_ms: attempt_duration_ms,
+            config: config
+          )
+        end
+      end
+
+      # Run a block in an OS thread from a shared pool, yielding the current fiber
+      # until it completes. Uses IO.pipe to yield the fiber to the Async event loop
+      # while the thread does work.
+      #
+      # The action is wrapped with a cancellation flag so that if the invocation
+      # finishes (e.g., suspended, terminal error) before the pool picks up the job,
+      # the action is skipped.
+      #
+      # Note: With Async 2.x and Ruby 3.1+, the Fiber Scheduler already intercepts
+      # most blocking I/O (Net::HTTP, TCPSocket, etc.) and yields the fiber
+      # automatically. +background: true+ is only needed for CPU-heavy native
+      # extensions that release the GVL (e.g., image processing, crypto).
+      def offload_to_thread(action)
+        read_io, write_io = IO.pipe
+        result = nil
+        error = nil
+        event = @attempt_finished_event
+
+        begin
+          BackgroundPool.submit do
+            if event.set?
+              # Attempt already finished before pool picked up the job — skip.
+              next
+            end
+
+            result = action.call
+          rescue Exception => e # rubocop:disable Lint/RescueException
+            error = e
+          ensure
+            write_io.close unless write_io.closed?
+          end
+
+          # Yields the fiber in Async context; resumes when the worker closes write_io.
+          read_io.read(1)
+          read_io.close
+
+          raise error if error
+
+          result
+        ensure
+          read_io.close unless read_io.closed?
+          write_io.close unless write_io.closed?
+        end
+      end
+
+      # A simple fixed-size thread pool for background: true runs.
+      # Avoids creating a new Thread per call (~1ms + ~1MB stack each).
+      # Workers are daemon threads that do not prevent process exit.
+      module BackgroundPool
+        @queue = Queue.new
+        @workers = []
+        @mutex = Mutex.new
+        @size = 0
+
+        POOL_SIZE = Integer(ENV.fetch('RESTATE_BACKGROUND_POOL_SIZE', 8))
+
+        module_function
+
+        # Submit a block to be executed by a pool worker.
+        def submit(&block)
+          ensure_started
+          @queue.push(block)
+        end
+
+        def ensure_started
+          return if @size >= POOL_SIZE
+
+          @mutex.synchronize do
+            while @size < POOL_SIZE
+              @size += 1
+              worker = Thread.new do
+                Kernel.loop do
+                  job = @queue.pop
+                  break if job == :shutdown
+
+                  job.call
+                end
+              end
+              worker.name = "restate-bg-#{@size}"
+              # Daemon thread: does not prevent the process from exiting.
+              worker.report_on_exception = false
+              @workers << worker
+            end
+          end
+        end
+      end
+    end
+  end
+end