streamdal 0.0.1

data/lib/streamdal.rb

@@ -0,0 +1,852 @@
+# frozen_string_literal: true
+require 'base64'
+require 'httparty'
+require 'logger'
+require 'securerandom'
+require 'wasmtime'
+require 'sp_sdk_pb'
+require 'sp_common_pb'
+require 'sp_info_pb'
+require 'sp_internal_pb'
+require "sp_internal_services_pb"
+require 'sp_pipeline_pb'
+require "sp_wsm_pb"
+require "steps/sp_steps_httprequest_pb"
+require "steps/sp_steps_kv_pb"
+require 'timeout'
+require 'google/protobuf'
+require_relative 'audiences'
+require_relative 'hostfunc'
+require_relative 'kv'
+require_relative 'metrics'
+require_relative 'schema'
+require_relative 'tail'
+require_relative 'validation'
+
+DEFAULT_GRPC_RECONNECT_INTERVAL = 5 # 5 seconds
+DEFAULT_PIPELINE_TIMEOUT = 1 / 10 # 100 milliseconds
+DEFAULT_STEP_TIMEOUT = 1 / 100 # 10 milliseconds
+DEFAULT_GRPC_TIMEOUT = 5 # 5 seconds
+DEFAULT_HEARTBEAT_INTERVAL = 1 # 1 second
+MAX_PAYLOAD_SIZE = 1024 * 1024 # 1 megabyte
+
+module Streamdal
+
+  OPERATION_TYPE_PRODUCER = 2
+  OPERATION_TYPE_CONSUMER = 1
+  CLIENT_TYPE_SDK = 1
+  CLIENT_TYPE_SHIM = 2
+
+  # Data class to hold instantiated wasm functions
+  class WasmFunction
+
+    ##
+    # Instance of an initialized wasm module and associated memory store
+
+    attr_accessor :instance, :store
+
+    def initialize
+      @instance = nil
+      @store = nil
+    end
+
+  end
+
+  ##
+  # Data class to store/pass audiences
+
+  Audience = Struct.new(:operation_type, :operation_name, :component_name) do
+    def to_proto(service_name)
+      Streamdal::Protos::Audience.new(
+        operation_type: Streamdal::Protos::OperationType.lookup(operation_type.to_i),
+        operation_name: operation_name,
+        component_name: component_name,
+        service_name: service_name,
+      )
+    end
+  end
+
+
+  class Client
+
+    ##
+    # Streamdal SDK Client
+    #
+    # There is only one public method: process(data, audience)
+    include Audiences
+    include Validation
+    include Schemas
+
+    # Aliases to keep lines short
+    CounterEntry = Streamdal::Metrics::CounterEntry
+    Metrics = Streamdal::Metrics
+
+    def initialize(cfg)
+      _validate_cfg(cfg)
+
+      @cfg = cfg
+      @log = cfg[:log]
+      @functions = {}
+      @session_id = SecureRandom.uuid
+      @pipelines = {}
+      @audiences = {}
+      @schemas = {}
+      @tails = {}
+      @paused_tails = {}
+      @metrics = Metrics.new(cfg)
+      @workers = []
+      @exit = false
+      @kv = Streamdal::KeyValue.new
+      @hostfunc = Streamdal::HostFunc.new(@kv)
+
+      # # Connect to Streamdal External gRPC API
+      @stub = Streamdal::Protos::Internal::Stub.new(@cfg[:streamdal_url], :this_channel_is_insecure)
+
+      _pull_initial_pipelines
+
+      @workers << Thread.new { _heartbeat }
+      @workers << Thread.new { _register }
+    end
+
+    def shutdown
+      # Set exit flag so workers exit
+      @exit = true
+
+      # Let loops exit
+      sleep(1)
+
+      # Exit any remaining threads
+      @workers.each do |w|
+        if w.running?
+          w.exit
+        end
+      end
+    end
+
+    def process(data, audience)
+      if data.length == 0
+        raise "data is required"
+      end
+
+      if audience.nil?
+        raise "audience is required"
+      end
+
+      resp = Streamdal::Protos::SDKResponse.new
+      resp.status = :EXEC_STATUS_TRUE
+      resp.pipeline_status = Google::Protobuf::RepeatedField.new(:message, Streamdal::Protos::PipelineStatus, [])
+      resp.data = data
+
+      aud = audience.to_proto(@cfg[:service_name])
+
+      labels = {
+        "service": @cfg[:service_name],
+        "operation_type": aud.operation_type,
+        "operation": aud.operation_name,
+        "component": aud.component_name,
+        "pipeline_name": "",
+        "pipeline_id": "",
+      }
+
+      # TODO: metrics
+      bytes_processed = Metrics::COUNTER_CONSUME_BYTES
+      errors_counter = Metrics::COUNTER_CONSUME_ERRORS
+      total_counter = Metrics::COUNTER_CONSUME_PROCESSED
+      rate_processed = Metrics::COUNTER_CONSUME_PROCESSED_RATE
+
+      if aud.operation_type == OPERATION_TYPE_PRODUCER
+        bytes_processed = Metrics::COUNTER_PRODUCE_BYTES
+        errors_counter = Metrics::COUNTER_PRODUCE_ERRORS
+        total_counter = Metrics::COUNTER_PRODUCE_PROCESSED
+        rate_processed = Metrics::COUNTER_PRODUCE_PROCESSED_RATE
+      end
+
+      payload_size = data.length
+
+      if payload_size > MAX_PAYLOAD_SIZE
+        # TODO: add metrics
+        resp.status = :EXEC_STATUS_ERROR
+        resp.error = "payload size exceeds maximum allowed size"
+        resp
+      end
+
+      # Needed for send_tail()
+      original_data = data
+
+      pipelines = _get_pipelines(aud)
+      if pipelines.length == 0
+        _send_tail(aud, "", original_data, original_data)
+        return resp
+      end
+
+      @metrics.incr(CounterEntry.new(bytes_processed, aud, labels, data.length))
+      @metrics.incr(CounterEntry.new(rate_processed, aud, labels, 1))
+
+      # Used for passing data between steps
+      isr = nil
+
+      Timeout::timeout(@cfg[:pipeline_timeout]) do
+        pipelines.each do |pipeline|
+          pipeline_status = Streamdal::Protos::PipelineStatus.new
+          pipeline_status.id = pipeline.id
+          pipeline_status.name = pipeline.name
+          pipeline_status.step_status = Google::Protobuf::RepeatedField.new(:message, Streamdal::Protos::StepStatus, [])
+
+          @log.debug "Running pipeline: '#{pipeline.name}'"
+
+          labels[:pipeline_id] = pipeline.id
+          labels[:pipeline_name] = pipeline.name
+
+          @metrics.incr(CounterEntry.new(total_counter, aud, labels, 1))
+          @metrics.incr(CounterEntry.new(bytes_processed, aud, labels, data.length))
+
+          pipeline.steps.each do |step|
+            step_status = Streamdal::Protos::StepStatus.new
+            step_status.name = step.name
+            step_status.status = :EXEC_STATUS_TRUE
+
+            begin
+              wasm_resp = _call_wasm(step, data, isr)
+            rescue => e
+              @log.error "Error running step '#{step.name}': #{e}"
+              step_status.status = :EXEC_STATUS_ERROR
+              step_status.error = e.to_s
+              pipeline_status.step_status.push(step_status)
+              break
+            end
+
+            if @cfg[:dry_run]
+              @log.debug "Running step '#{step.name}' in dry-run mode"
+            end
+
+            if wasm_resp.output_payload.length > 0
+              resp.data = wasm_resp.output_payload
+            end
+
+            _handle_schema(aud, step, wasm_resp)
+
+            isr = wasm_resp.inter_step_result
+
+            case wasm_resp.exit_code
+            when :WASM_EXIT_CODE_FALSE
+              cond = step.on_false
+              cond_type = :CONDITION_TYPE_ON_FALSE
+              exec_status = :EXEC_STATUS_FALSE
+            when :WASM_EXIT_CODE_ERROR
+              cond = step.on_error
+              cond_type = :CONDITION_TYPE_ON_ERROR
+              exec_status = :EXEC_STATUS_ERROR
+              isr = nil
+
+              # errors_counter, 1, labels, aud
+              @metrics.incr(CounterEntry.new(errors_counter, aud, labels, 1))
+            else
+              cond = step.on_true
+              exec_status = :EXEC_STATUS_TRUE
+              cond_type = :CONDITION_TYPE_ON_TRUE
+            end
+
+            _notify_condition(pipeline, step, aud, cond, resp.data, cond_type)
+
+            if @cfg[:dry_run]
+              @log.debug("Step '#{step.name}' completed with status: #{exec_status}, continuing to next step")
+              next
+            end
+
+            # Whether we are aborting early, aborting current, or continuing, we need to set the step status
+            step_status.status = exec_status
+            step_status.status_message = "Step returned: #{wasm_resp.exit_msg}"
+
+            # Pull metadata from step into SDKResponse
+            unless cond.nil?
+              resp.metadata = cond.metadata
+
+              case cond.abort
+              when :ABORT_CONDITION_ABORT_CURRENT
+                step_status.status = exec_status
+                step_status.status_message = "Step returned: #{wasm_resp.exit_msg}"
+                pipeline_status.step_status.push(step_status)
+                resp.pipeline_status.push(pipeline_status)
+                # Continue outer pipeline loop, there might be additional pipelines
+                break
+              when :ABORT_CONDITION_ABORT_ALL
+                # Set step status and push to pipeline status
+                step_status.status = exec_status
+                step_status.status_message = "Step returned: #{wasm_resp.exit_msg}"
+                pipeline_status.step_status.push(step_status)
+                resp.pipeline_status.push(pipeline_status)
+
+                # Since we're returning early here, also need to set the response status
+                resp.status = exec_status
+                resp.status_message = "Step returned: #{wasm_resp.exit_msg}"
+
+                _send_tail(aud, pipeline.id, original_data, resp.data)
+                return resp
+              else
+                # Do nothing
+              end
+            end
+
+            # Append step status to the current pipeline status' array
+            pipeline_status.step_status.push(step_status)
+          end
+
+          # Append pipeline status to the response
+          resp.pipeline_status.push(pipeline_status)
+        end # pipelines.each
+      end # timeout
+
+      _send_tail(aud, "", original_data, resp.data)
+
+      if @cfg[:dry_run]
+        @log.debug "Dry-run, setting response data to original data"
+        resp.data = original_data
+      end
+
+      resp
+    end
+
+    private
+
+    def _validate_cfg(cfg)
+      if cfg[:streamdal_url].nil? || cfg[:streamdal_url].empty?
+        raise "streamdal_url is required"
+      end
+
+      if cfg[:streamdal_token].nil? || cfg[:streamdal_token].empty?
+        raise "streamdal_token is required"
+      end
+
+      if cfg[:service_name].nil? || cfg[:streamdal_token].empty?
+        raise "service_name is required"
+      end
+
+      if cfg[:log].nil? || cfg[:streamdal_token].empty?
+        logger = Logger.new(STDOUT)
+        logger.level = Logger::ERROR
+        cfg[:log] = logger
+      end
+
+      if cfg[:pipeline_timeout].nil?
+        cfg[:pipeline_timeout] = DEFAULT_PIPELINE_TIMEOUT
+      end
+
+      if cfg[:step_timeout].nil?
+        cfg[:step_timeout] = DEFAULT_STEP_TIMEOUT
+      end
+    end
+
+    def _handle_command(cmd)
+      case cmd.command.to_s
+      when "kv"
+        _handle_kv(cmd)
+      when "tail"
+        _handle_tail_request(cmd)
+      when "set_pipelines"
+        _set_pipelines(cmd)
+      when "keep_alive"
+        # Do nothing
+      else
+        @log.error "unknown command type #{cmd.command}"
+      end
+    end
+
+    def _handle_kv(cmd)
+      begin
+        validate_kv_command(cmd)
+      rescue => e
+        @log.error "KV command validation failed: #{e}"
+        return nil
+      end
+
+      cmd.kv.instructions.each do |inst|
+        validate_kv_instruction(inst)
+
+        case inst.action
+        when :KV_ACTION_CREATE
+          @kv.set(inst.key, inst.value)
+        when :KV_ACTION_UPDATE
+          @kv.set(inst.key, inst.value)
+        when :KV_ACTION_DELETE
+          @kv.delete(inst.key)
+        when :KV_ACTION_DELETE_ALL
+          @kv.purge
+        else
+          @log.error "Unknown KV action: '#{inst.action}'"
+        end
+      end
+    end
+
+    def _set_pipelines(cmd)
+      if cmd.nil?
+        raise "cmd is required"
+      end
+
+      cmd.set_pipelines.pipelines.each_with_index { |p, pIdx|
+        p.steps.each_with_index { |step, idx|
+          if step._wasm_bytes == ""
+            if cmd.set_pipelines.wasm_modules.has_key?(step._wasm_id)
+              step._wasm_bytes = cmd.set_pipelines.wasm_modules[step._wasm_id].bytes
+              cmd.set_pipelines.pipelines[pIdx].steps[idx] = step
+            else
+              @log.error "WASM module not found for step: #{step._wasm_id}"
+            end
+          end
+        }
+
+        aud_str = aud_to_str(cmd.audience)
+        @pipelines.key?(aud_str) ? @pipelines[aud_str].push(p) : @pipelines[aud_str] = [p]
+      }
+    end
+
+    def _pull_initial_pipelines
+      req = Streamdal::Protos::GetSetPipelinesCommandsByServiceRequest.new
+      req.service_name = @cfg[:service_name]
+
+      resp = @stub.get_set_pipelines_commands_by_service(req, metadata: _metadata)
+
+      @log.debug "Received '#{resp.set_pipeline_commands.length}' initial pipelines"
+
+      resp.set_pipeline_commands.each do |cmd|
+        cmd.set_pipelines.wasm_modules = resp.wasm_modules
+        _set_pipelines(cmd)
+      end
+    end
+
+    def _get_function(step)
+      # We cache functions so we can eliminate the wasm bytes from steps to save on memory
+      # And also to avoid re-initializing the same function multiple times
+      if @functions.key?(step._wasm_id)
+        return @functions[step._wasm_id]
+      end
+
+      engine = Wasmtime::Engine.new
+      mod = Wasmtime::Module.new(engine, step._wasm_bytes)
+      linker = Wasmtime::Linker.new(engine, wasi: true)
+
+      linker.func_new("env", "httpRequest", [:i32, :i32], [:i64]) do |caller, ptr, len|
+        @hostfunc.http_request(caller, ptr, len)
+      end
+
+      linker.func_new("env", "kvExists", [:i32, :i32], [:i64]) do |caller, ptr, len|
+        @hostfunc.kv_exists(caller, ptr, len)
+      end
+
+      wasi_ctx = Wasmtime::WasiCtxBuilder.new
+                                         .inherit_stdout
+                                         .inherit_stderr
+                                         .set_argv(ARGV)
+                                         .set_env(ENV)
+                                         .build
+      store = Wasmtime::Store.new(engine, wasi_ctx: wasi_ctx)
+
+      instance = linker.instantiate(store, mod)
+
+      # TODO: host funcs
+
+      # Store in cache
+      func = WasmFunction.new
+      func.instance = instance
+      func.store = store
+      @functions[step._wasm_id] = func
+
+      func
+    end
+
+    def _call_wasm(step, data, isr)
+      if step.nil?
+        raise "step is required"
+      end
+
+      if data.nil?
+        raise "data is required"
+      end
+
+      if isr.nil?
+        isr = Streamdal::Protos::InterStepResult.new
+      end
+
+      req = Streamdal::Protos::WASMRequest.new
+      req.step = step.clone
+      req.input_payload = data
+      req.inter_step_result = isr
+
+      begin
+        Timeout::timeout(@cfg[:step_timeout]) do
+          wasm_resp = _exec_wasm(req)
+          return Streamdal::Protos::WASMResponse.decode(wasm_resp)
+        end
+      rescue => e
+        resp = Streamdal::Protos::WASMResponse.new
+        resp.exit_code = :WASM_EXIT_CODE_ERROR
+        resp.exit_msg = "Failed to execute WASM: #{e}"
+        resp.output_payload = ""
+        return resp
+      end
+    end
+
+    def _gen_register_request
+      req = Streamdal::Protos::RegisterRequest.new
+      req.service_name = @cfg[:service_name]
+      req.session_id = @session_id
+      req.dry_run = @cfg[:dry_run] || false
+      req.client_info = _gen_client_info
+
+      req
+    end
+
+    def _gen_client_info
+      arch, os = RUBY_PLATFORM.split(/-/)
+
+      ci = Streamdal::Protos::ClientInfo.new
+      ci.client_type = :CLIENT_TYPE_SDK
+      ci.library_name = "ruby-sdk"
+      ci.library_version = "0.0.1"
+      ci.language = "ruby"
+      ci.arch = arch
+      ci.os = os
+
+      ci
+    end
+
+    # Returns metadata for gRPC requests to the internal gRPC API
+    def _metadata
+      { "auth-token" => @cfg[:streamdal_token].to_s }
+    end
+
+    def _register
+      @log.info("register started")
+
+      # Register with Streamdal External gRPC API
+      resps = @stub.register(_gen_register_request, metadata: _metadata)
+      resps.each do |r|
+        if @exit
+          break
+        end
+
+        _handle_command(r)
+      end
+
+      @log.info("register exited")
+    end
+
+    def _exec_wasm(req)
+      wasm_func = _get_function(req.step)
+
+      # Empty out _wasm_bytes, we don't need it anymore
+      # TODO: does this actually update the original object?
+      req.step._wasm_bytes = ""
+
+      data = req.to_proto
+
+      memory = wasm_func.instance.export("memory").to_memory
+      alloc = wasm_func.instance.export("alloc").to_func
+      dealloc = wasm_func.instance.export("dealloc").to_func
+      f = wasm_func.instance.export("f").to_func
+
+      start_ptr = alloc.call(data.length)
+
+      memory.write(start_ptr, data)
+
+      # Result is a 64bit int where the first 32 bits are the pointer to the result
+      # and the last 32 bits are the length of the result. This is due to the fact
+      # that we can only return an integer from a wasm function.
+      result_ptr = f.call(start_ptr, data.length)
+      ptr_true = result_ptr >> 32
+      len_true = result_ptr & 0xFFFFFFFF
+
+      res = memory.read(ptr_true, len_true)
+
+      # Dealloc result memory since we already read it
+      dealloc.call(ptr_true, res.length)
+
+      res
+    end
+
+    def _get_pipelines(aud)
+      aud_str = aud_to_str(aud)
+
+      _add_audience(aud)
+
+      if @pipelines.key?(aud_str)
+        return @pipelines[aud_str]
+      end
+
+      []
+    end
+
+    def _heartbeat
+      until @exit
+        req = Streamdal::Protos::HeartbeatRequest.new
+        req.session_id = @session_id
+        req.audiences = Google::Protobuf::RepeatedField.new(:message, Streamdal::Protos::Audience, [])
+
+        @audiences.each do |_, aud|
+          req.audiences.push(aud)
+        end
+
+        req.client_info = _gen_client_info
+        req.service_name = @cfg[:service_name]
+
+        @stub.heartbeat(req, metadata: _metadata)
+        sleep(DEFAULT_HEARTBEAT_INTERVAL)
+      end
+    end
+
+    ######################################################################################
+    # Tail methods
+    ######################################################################################
+
+    def _handle_tail_request(cmd)
+      validate_tail_request(cmd)
+
+      case cmd.tail.request.type
+      when :TAIL_REQUEST_TYPE_START
+        _start_tail(cmd)
+      when :TAIL_REQUEST_TYPE_STOP
+        _stop_tail(cmd)
+      when :TAIL_REQUEST_TYPE_PAUSE
+        _pause_tail(cmd)
+      when :TAIL_REQUEST_TYPE_RESUME
+        _resume_tail(cmd)
+      else
+        raise "unknown tail request type: '#{cmd.tail.request.type.inspect}'"
+      end
+    end
+
+    def _get_active_tails_for_audience(aud)
+      aud_str = aud_to_str(aud)
+      if @tails.key?(aud_str)
+        return @tails[aud_str].values
+      end
+
+      []
+    end
+
+    def _send_tail(aud, pipeline_id, original_data, new_data)
+      tails = _get_active_tails_for_audience(aud)
+      if tails.length == 0
+        return nil
+      end
+
+      tails.each do |tail|
+        req = Streamdal::Protos::TailResponse.new
+        req.type = :TAIL_RESPONSE_TYPE_PAYLOAD
+        req.audience = aud
+        req.pipeline_id = pipeline_id
+        req.session_id = @session_id
+        req.timestamp_ns = Time.now.to_i
+        req.original_data = original_data
+        req.new_data = new_data
+        req.tail_request_id = tail.request.id
+        tail.queue.push(req)
+      end
+    end
+
+    def _notify_condition(pipeline, step, aud, cond, data, cond_type)
+      if cond.nil?
+        return nil
+      end
+
+      if cond.notification.nil?
+        return nil
+      end
+
+      @log.debug "Notifying"
+
+      if @cfg[:dry_run]
+        return nil
+      end
+
+      @metrics.incr(CounterEntry.new(Metrics::COUNTER_NOTIFY, aud, {
+        "service": @cfg[:service_name],
+        "component_name": aud.component_name,
+        "pipeline_name": pipeline.name,
+        "pipeline_id": pipeline.id,
+        "operation_name": aud.operation_name,
+      }, 1))
+
+      req = Streamdal::Protos::NotifyRequest.new
+      req.audience = aud
+      req.pipeline_id = pipeline.id
+      req.step = step
+      req.payload = data
+      req.condition_type = cond_type
+      req.occurred_at_unix_ts_utc = Time.now.to_i
+
+      Thread.new do
+        @stub.notify(req, metadata: _metadata)
+      end
+    end
+
+    def _start_tail(cmd)
+      validate_tail_request(cmd)
+
+      req = cmd.tail.request
+      @log.debug "Starting tail '#{req.id}'"
+
+      aud_str = aud_to_str(cmd.tail.request.audience)
+
+      # Do we already have a tail for this audience
+      if @tails.key?(aud_str) && @tails[aud_str].key?(req.id)
+        @log.error "Tail '#{req.id}' already exists, skipping TailCommand"
+        return
+      end
+
+      @log.debug "Tailing audience: #{aud_str}"
+
+      t = Streamdal::Tail.new(
+        req,
+        @cfg[:streamdal_url],
+        @cfg[:streamdal_token],
+        @cfg[:log],
+        @metrics
+      )
+
+      t.start_tail_workers
+
+      _set_active_tail(t)
+
+    end
+
+    def _set_active_tail(tail)
+      key = aud_to_str(tail.request.audience)
+
+      unless @tails.key?(key)
+        @tails[key] = {}
+      end
+
+      @tails[key][tail.request.id] = tail
+    end
+
+    def _set_paused_tail(tail)
+      key = aud_to_str(tail.request.aud)
+
+      unless @paused_tails.key?(key)
+        @paused_tails[key] = {}
+      end
+
+      @paused_tails[key][tail.request.id] = tail
+    end
+
+    def _stop_tail(cmd)
+      @log.debug "Stopping tail '#{cmd.tail.request.id}'"
+      key = aud_to_str(cmd.tail.request.audience)
+
+      if @tails.key?(key) && @tails[key].key?(cmd.tail.request.id)
+        @tails[key][cmd.tail.request.id].stop_tail
+
+        # Remove from active tails
+        @tails[key].delete(cmd.tail.request.id)
+
+        if @tails[key].length == 0
+          @tails.delete(key)
+        end
+      end
+
+      if @paused_tails.key?(key) && @paused_tails[key].key?(cmd.tail.request.id)
+        @paused_tails[key].delete(cmd.tail.request.id)
+
+        if @paused_tails[key].length == 0
+          @paused_tails.delete(key)
+        end
+      end
+    end
+
+    def _stop_all_tails
+      # TODO: does this modify the instances variables or copy them?
+      _stop_tails(@tails)
+      _stop_tails(@paused_tails)
+    end
+
+    def _stop_tails(tails = {})
+      # Helper method for _stop_all_tails
+      tails.each do |aud, aud_tails|
+        aud_tails.each do |t|
+          t.stop_tail
+          tails[aud].delete(tail.request.id)
+
+          if tails[aud].length == 0
+            tails.delete(aud)
+          end
+        end
+      end
+    end
+
+    def _pause_tail(cmd)
+      t = _remove_active_tail(cmd.tail.request.audience, cmd.tail.request.tail.id)
+      t.stop_tail
+
+      _set_paused_tail(t)
+
+      @log.debug "Paused tail '#{cmd.tail.request.tail.id}'"
+    end
+
+    def _resume_tail(cmd)
+      t = _remove_paused_tail(cmd.tail.request.audience, cmd.tail.request.tail.id)
+      if t.nil?
+        @log.error "Tail '#{cmd.tail.request.tail.id}' not found in paused tails"
+        return nil
+      end
+
+      t.start_tail_workers
+
+      _set_active_tail(t)
+
+      @log.debug "Resumed tail '#{cmd.tail.request.tail.id}'"
+    end
+
+    def _remove_active_tail(aud, tail_id)
+      key = aud_to_str(aud)
+
+      if @tails.key?(key) && @tails[key].key?(tail_id)
+        t = @tails[key][tail_id]
+        t.stop_tail
+
+        @tails[key].delete(tail_id)
+
+        if @tails[key].length == 0
+          @tails.delete(key)
+        end
+
+        t
+      end
+    end
+
+    def _remove_paused_tail(aud, tail_id)
+      key = aud_to_str(aud)
+
+      if @paused_tails.key?(key) && @paused_tails[key].key?(tail_id)
+        t = @paused_tails[key][tail_id]
+
+        @paused_tails[key].delete(tail_id)
+
+        if @paused_tails[key].length == 0
+          @paused_tails.delete(key)
+        end
+
+        t
+      end
+    end
+
+    # Called by host functions to write memory to wasm instance so that
+    # the wasm module can read the result of a host function call
+    def write_to_memory(caller, res)
+      alloc = caller.export("alloc").to_func
+      memory = caller.export("memory").to_memory
+
+      # Serialize protobuf message
+      resp = res.to_proto
+
+      # Allocate memory for response
+      resp_ptr = alloc.call(resp.length)
+
+      # Write response to memory
+      memory.write(resp_ptr, resp)
+
+      # return 64bit integer where first 32 bits is the pointer, and the last 32 is the length
+      resp_ptr << 32 | resp.length
+    end
+
+  end
+end