streamdal 0.0.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: cdf556c16ea4239101cc21f803b78c458c5e2490a1c5df6b2211f5268cc94799
+  data.tar.gz: 1bca44db3048358a5bdf978504a8559c1a81ede836a11269c5ab8020698d71b6
+SHA512:
+  metadata.gz: 9f0fb7a83fc6e924f520d5402385bd857aff3a27884fd2a1aa9572a002723ad26d3849a2930829e6a05e450547ab1e80293992a693bbeb2bb9bacad8756011f9
+  data.tar.gz: 2176e44cd407be37c0137f48c77f7a56550cddacc36a5a8043643049b416138c1b96b35c850c630dfd1f9febdc4e72e9224fb055b8dfd871c54e9bef5fbdbe47

data/lib/audiences.rb

@@ -0,0 +1,45 @@
+module Audiences
+  def aud_to_str(aud)
+    "#{aud.service_name}.#{aud.component_name}.#{aud.operation_type}.#{aud.operation_name}"
+  end
+
+  def str_to_aud(str)
+    # TODO: move to common package
+    parts = str.split(".")
+    aud = Streamdal::Protos::Audience.new
+    aud.service_name = parts[0]
+    aud.component_name = parts[1]
+    aud.operation_type = parts[2]
+    aud.operation_name = parts[3]
+    aud
+  end
+
+  def _seen_audience(aud)
+    @audiences.key?(aud_to_str(aud))
+  end
+
+  def _add_audience(aud)
+    # Add an audience to the local cache map and send to server
+    if _seen_audience(aud)
+      return
+    end
+
+    @audiences[aud_to_str(aud)] = aud
+
+    req = Streamdal::Protos::NewAudienceRequest.new
+    req.session_id = @session_id
+    req.audience = aud
+    @stub.new_audience(req, metadata: _metadata)
+  end
+
+  def _add_audiences
+    # This method is used to re-announce audiences after a disconnect
+
+    @audiences.each do |aud|
+      req = Streamdal::Protos::NewAudienceRequest.new
+      req.session_id = @session_id
+      req.audience = aud
+      @stub.new_audience(req, metadata: _metadata)
+    end
+  end
+end

data/lib/audiences_spec.rb

@@ -0,0 +1,5 @@
+require_relative 'spec_helper'
+
+RSpec.describe "Streamdal::Audiences" do
+
+end

data/lib/hostfunc.rb

@@ -0,0 +1,109 @@
+require "steps/sp_steps_kv_pb"
+
+module Streamdal
+  class HostFunc
+
+    ##
+    # This class holds methods that are called by wasm modules
+
+    def initialize(kv)
+      @kv = kv
+    end
+
+    ##
+    # kv_exists is a host function that is used to check if a key exists in the KV store
+    def kv_exists(caller, ptr, len)
+
+      data = caller.export("memory").to_memory.read(ptr, len)
+
+      # Read request from memory and decode into HttpRequest
+      req = Streamdal::Protos::KVStep.decode(data)
+
+      exists = @kv.exists(req.key)
+
+      msg = exists ? "Key '#{req.key}' exists" : "Key #{req.key} does not exist"
+
+      status = exists ? :KV_STATUS_SUCCESS : :KV_STATUS_FAILURE
+
+      wasm_resp = Streamdal::Protos::KVStepResponse.new
+      wasm_resp.status = status
+      wasm_resp.message = msg
+
+      write_to_memory(caller, wasm_resp)
+    end
+
+    ##
+    # http_request performs a http request on behalf of a wasm module since WASI cannot talk sockets
+    def http_request(caller, ptr, len)
+      data = caller.export("memory").to_memory.read(ptr, len)
+
+      # Read request from memory and decode into HttpRequest
+      req = Streamdal::Protos::HttpRequest.decode(data)
+
+      # Attempt to make HTTP request
+      # On error, return a mock 400 response with the error as the body
+      begin
+        response = _make_http_request(req)
+      rescue => e
+        wasm_resp = Streamdal::Protos::HttpResponse.new
+        wasm_resp.code = 400
+        wasm_resp.body = "Unable to execute HTTP request: #{e}"
+        return wasm_resp
+      end
+
+      # Successful request, build the proto from the httparty response
+      wasm_resp = Streamdal::Protos::HttpResponse.new
+      wasm_resp.code = response.code
+      wasm_resp.body = response.body
+      wasm_resp.headers = Google::Protobuf::Map.new(:string, :string, {})
+
+      # Headers can have multiple values, but we just want a map[string]string here for simplicity
+      # The client can pase by the delimiter ";" if needed.
+      response.headers.each do |k, values|
+        wasm_resp.headers[k] = values.kind_of?(Array) ? values.join("; ") : values
+      end
+
+      # Write the HttpResponse proto message to WASM memory
+      # The .wasm module will read/decode this data internally
+      write_to_memory(caller, wasm_resp)
+    end
+
+    private
+
+    ##
+    # Performs an http request
+    def _make_http_request(req)
+      if req.nil?
+        raise "req is required"
+      end
+
+      options = {
+        headers: { "Content-Type": "application/json", },
+      }
+
+      req.headers.each { |key, value| options.headers[key] = value }
+
+      case req.to_h[:method]
+      when :HTTP_REQUEST_METHOD_GET
+        return HTTParty.get(req.url)
+      when :HTTP_REQUEST_METHOD_POST
+        options.body = req.body
+        return HTTParty.post(req.url, options)
+      when :HTTP_REQUEST_METHOD_PUT
+        options.body = req.body
+        return HTTParty.put(req.url, options)
+      when :HTTP_REQUEST_METHOD_DELETE
+        return HTTParty.delete(req.url)
+      when :HTTP_REQUEST_METHOD_PATCH
+        options.body = req.body
+        return HTTParty.patch(req.url, options)
+      when :HTTP_REQUEST_METHOD_HEAD
+        return HTTParty.head(req.url)
+      when :HTTP_REQUEST_METHOD_OPTIONS
+        return HTTParty.options(req.url)
+      else
+        raise ArgumentError, "Invalid http request method: #{req.method}"
+      end
+    end
+  end
+end

data/lib/hostfunc_spec.rb

@@ -0,0 +1,5 @@
+require_relative 'spec_helper'
+
+RSpec.describe "Streamdal::HostFunc" do
+
+end

data/lib/kv.rb

@@ -0,0 +1,52 @@
+module Streamdal
+  class KeyValue
+    def initialize
+      @kvs = {}
+      @mtx = Mutex.new
+    end
+
+    def set(key, value)
+      @mtx.synchronize do
+        @kvs[key] = value
+      end
+    end
+
+    def get(key)
+      @mtx.synchronize do
+        @kvs[key]
+      end
+    end
+
+    def delete(key)
+      @mtx.synchronize do
+        @kvs.delete(key)
+      end
+    end
+
+    def keys
+      @mtx.synchronize do
+        @kvs.keys
+      end
+    end
+
+    def items
+      @mtx.synchronize do
+        @kvs.values
+      end
+    end
+
+    def exists(key)
+      @mtx.synchronize do
+        @kvs.key?(key)
+      end
+    end
+
+    def purge
+      @mtx.synchronize do
+        num_keys = @kvs.keys.length
+        @kvs = {}
+        num_keys
+      end
+    end
+  end
+end

data/lib/kv_spec.rb

@@ -0,0 +1,54 @@
+require 'rspec'
+require_relative 'spec_helper'
+require_relative 'kv'
+
+RSpec.describe 'KeyValue' do
+  before(:each) do
+    @kv = Streamdal::KeyValue.new
+  end
+
+  it 'stores and retrieves key-value pair' do
+    @kv.set('key', 'value')
+    expect(@kv.get('key')).to eq('value')
+  end
+
+  it 'deletes a key' do
+    expect(@kv.keys.length).to eq(0)
+    @kv.set('key1', 'value')
+    @kv.set('key2', 'value')
+    expect(@kv.keys.length).to eq(2)
+    @kv.delete('key1')
+    expect(@kv.keys.length).to eq(1)
+  end
+
+  it 'returns array of keys' do
+    expect(@kv.keys.length).to eq(0)
+    @kv.set('key1', 'value')
+    expect(@kv.keys.length).to eq(1)
+    expect(@kv.keys[0]).to eq('key1')
+  end
+
+  it 'returns array of item values' do
+    expect(@kv.keys.length).to eq(0)
+    @kv.set('key1', 'value1')
+    expect(@kv.keys.length).to eq(1)
+    expect(@kv.items[0]).to eq('value1')
+  end
+
+  it 'returns if a key exists or not' do
+    expect(@kv.exists('key1')).to eq(false)
+    @kv.set('key1', 'value')
+    expect(@kv.exists('key1')).to eq(true)
+  end
+
+  it 'purges all keys' do
+    expect(@kv.keys.length).to eq(0)
+    @kv.set('key1', 'value')
+    @kv.set('key2', 'value')
+    expect(@kv.keys.length).to eq(2)
+    num_keys = @kv.purge
+    expect(@kv.keys.length).to eq(0)
+    expect(num_keys).to eq(2)
+  end
+
+end

data/lib/metrics.rb

@@ -0,0 +1,265 @@
+module Streamdal
+  class Counter
+    attr_accessor :last_updated, :name, :aud, :labels
+
+    def initialize(name, aud, labels = {}, value = 0.0)
+      @name = name
+      @aud = aud
+      @labels = labels
+      @value = 0
+      @last_updated = Time::now
+      @value_mtx = Mutex.new
+    end
+
+    def incr(val)
+      @value_mtx.synchronize do
+        @value = @value + val
+        @last_updated = Time::now
+      end
+    end
+
+    def reset
+      @value_mtx.synchronize do
+        @value = 0.0
+      end
+    end
+
+    def val
+      @value_mtx.synchronize do
+        @value
+      end
+    end
+  end
+
+  class Metrics
+
+    COUNTER_CONSUME_BYTES = "counter_consume_bytes"
+    COUNTER_CONSUME_PROCESSED = "counter_consume_processed"
+    COUNTER_CONSUME_ERRORS = "counter_consume_errors"
+    COUNTER_PRODUCE_BYTES = "counter_produce_bytes"
+    COUNTER_PRODUCE_PROCESSED = "counter_produce_processed"
+    COUNTER_PRODUCE_ERRORS = "counter_produce_errors"
+    COUNTER_NOTIFY = "counter_notify"
+    COUNTER_DROPPED_TAIL_MESSAGES = "counter_dropped_tail_messages"
+    COUNTER_CONSUME_BYTES_RATE = "counter_consume_bytes_rate"
+    COUNTER_PRODUCE_BYTES_RATE = "counter_produce_bytes_rate"
+    COUNTER_CONSUME_PROCESSED_RATE = "counter_consume_processed_rate"
+    COUNTER_PRODUCE_PROCESSED_RATE = "counter_produce_processed_rate"
+
+    WORKER_POOL_SIZE = 3
+    DEFAULT_COUNTER_REAPER_INTERVAL = 10
+    DEFAULT_COUNTER_TTL = 10
+    DEFAULT_COUNTER_PUBLISH_INTERVAL = 1
+
+    CounterEntry = Struct.new(:name, :aud, :labels, :value)
+
+    def initialize(cfg)
+      if cfg.nil?
+        raise ArgumentError, "cfg is nil"
+      end
+
+      @cfg = cfg
+      @log = cfg[:log]
+      @counters = {}
+      @counters_mtx = Mutex.new
+      @exit = false
+      @incr_queue = Queue.new
+      @publish_queue = Queue.new
+      @workers = []
+      @stub = Streamdal::Protos::Internal::Stub.new(@cfg[:streamdal_url], :this_channel_is_insecure)
+
+      _start
+    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 self.composite_id(counter_name, labels = {})
+      if labels.nil?
+        labels = {}
+      end
+      "#{counter_name}-#{labels.values.join("-")}".freeze
+    end
+
+    def get_counter(ce)
+      if ce.nil?
+        raise ArgumentError, "ce is nil"
+      end
+
+      k = Metrics::composite_id(ce.name, ce.labels)
+
+      @counters_mtx.synchronize do
+        if @counters.key?(k)
+          @counters[k]
+        end
+      end
+
+      # No counter exists, create a new one and return it
+      new_counter(ce)
+    end
+
+    def new_counter(ce)
+      c = Counter.new(ce.name, ce.aud, ce.labels, ce.value)
+
+      @counters_mtx.synchronize do
+        @counters[Metrics::composite_id(ce.name, ce.labels)] = c
+      end
+
+      c
+    end
+
+    def incr(ce)
+      c = get_counter(ce)
+
+      if c.nil?
+        new_counter(ce)
+        nil
+      end
+
+      @incr_queue.push(ce)
+    end
+
+    def remove_counter(name)
+      @counters_mtx.synchronize do
+        @counters.delete(name)
+      end
+    end
+
+    private
+
+    def _start
+      WORKER_POOL_SIZE.times do |i|
+        @workers << Thread.new { _run_incrementer_worker(i) }
+        @workers << Thread.new { _run_publisher_worker(i) }
+      end
+
+      @workers << Thread.new { _run_publisher }
+      @workers << Thread.new { _run_reaper }
+    end
+
+    def _publish_metrics(ce)
+      metric = Streamdal::Protos::Metric.new
+      metric.name = ce.name
+      metric.labels = Google::Protobuf::Map.new(:string, :string, ce.labels)
+      metric.value = ce.value
+      metric.audience = ce.aud
+
+      req = Streamdal::Protos::MetricsRequest.new
+      req.metrics = Google::Protobuf::RepeatedField.new(:message, Streamdal::Protos::Metric, [metric])
+
+      @log.debug("Published metric: #{ce.name} #{ce.labels} #{ce.value}")
+      @stub.metrics(req, metadata: _metadata)
+    end
+
+    def _run_publisher
+      # Background thread that reads values from counters, adds them to the publish queue, and then
+      # resets the counter's value back to zero
+      unless @exit
+        @log.debug("Starting publisher")
+
+        # Sleep on startup and then and between each loop run
+        sleep(DEFAULT_COUNTER_PUBLISH_INTERVAL)
+
+        # Get all counters
+        # Loop over each counter, get the value,
+        #   if value > 0, continue
+        #   if now() - last_updated > 10 seconds, remove counter
+        # Grab copy of counters
+        @counters_mtx.lock
+        new_counters = @counters.dup
+        @counters_mtx.unlock
+
+        new_counters.each do |_, counter|
+          if counter.val == 0
+            next
+          end
+
+          ce = CounterEntry.new(counter.name, counter.aud, counter.labels, counter.val)
+          counter.reset
+
+          @publish_queue.push(ce)
+        end
+      end
+    end
+
+    def _run_publisher_worker(worker_id)
+      @log.debug("Starting publisher worker '#{worker_id}'")
+
+      until @exit
+        ce = @incr_queue.pop
+        if ce.nil?
+          next
+        end
+        begin
+          _publish_metrics(ce)
+        rescue => e
+          @log.error("Failed to publish metrics: #{e}: #{ce.inspect}")
+        end
+      end
+
+      @log.debug("Exiting publisher worker '#{worker_id}'")
+    end
+
+    def _run_reaper
+      @log.debug("Starting reaper")
+
+      until @exit
+        # Sleep on startup and then and between each loop run
+        sleep(DEFAULT_COUNTER_REAPER_INTERVAL)
+
+        # Get all counters
+        # Loop over each counter, get the value,
+        #   if value > 0, continue
+        #   if now() - last_updated > 10 seconds, remove counter
+        # Grab copy of counters
+        @counters_mtx.synchronize do
+          @counters.each do |name, counter|
+            if counter.val > 0
+              next
+            end
+
+            if Time::now - counter.last_updated > DEFAULT_COUNTER_TTL
+              @log.debug("Reaping counter '#{name}'")
+              @counters.delete(name)
+            end
+          end
+        end
+      end
+
+      @log.debug("Exiting reaper")
+    end
+
+    def _run_incrementer_worker(worker_id)
+      @log.debug("Starting incrementer worker '#{worker_id}'")
+
+      until @exit
+        ce = @incr_queue.pop
+
+        next if ce.nil?
+
+        c = get_counter(ce)
+
+        c.incr(ce.value)
+      end
+
+      @log.debug("Exiting incrementer worker '#{worker_id}'")
+    end
+
+    # Returns metadata for gRPC requests to the internal gRPC API
+    def _metadata
+      { "auth-token" => @cfg[:streamdal_token].to_s }
+    end
+  end
+end

data/lib/metrics_spec.rb

@@ -0,0 +1,5 @@
+require_relative 'spec_helper'
+
+RSpec.describe "Streamdal::Metrics" do
+
+end

data/lib/schema.rb

@@ -0,0 +1,47 @@
+include Streamdal::Protos
+
+module Schemas
+  def _set_schema(aud, schema)
+    s = Streamdal::Protos::Schema.new
+    s.json_schema = schema
+    @schemas[aud_to_str(aud)] = s
+  end
+
+  def _get_schema(aud)
+    if @schemas.key?(aud_to_str(aud))
+      return @schemas[aud_to_str(aud)].json_schema
+    end
+
+    ""
+  end
+
+  def _handle_schema(aud, step, wasm_resp)
+    # Only handle schema steps
+    if step.infer_schema.nil?
+      return nil
+    end
+
+    # Only successful schema inferences
+    if wasm_resp.exit_code != :WASM_EXIT_CODE_TRUE
+      return nil
+    end
+
+    # If existing schema matches, do nothing
+    existing_schema = _get_schema(aud)
+    if existing_schema == wasm_resp.output_step
+      return nil
+    end
+
+    _set_schema(aud, wasm_resp.output_step)
+
+    req = Streamdal::Protos::SendSchemaRequest.new
+    req.audience = aud
+    req.schema = Streamdal::Protos::Schema.new
+    req.schema.json_schema = wasm_resp.output_step
+
+    # Run in thread so we don't block on gRPC call
+    Thread.new do
+      @stub.send_schema(req, metadata: _metadata)
+    end
+  end
+end

data/lib/schema_spec.rb

@@ -0,0 +1,59 @@
+require 'rspec'
+require 'sp_wsm_pb'
+require 'steps/sp_steps_inferschema_pb'
+require 'sp_pipeline_pb'
+require 'sp_common_pb'
+require_relative 'streamdal'
+require_relative 'spec_helper'
+require_relative 'schema'
+require_relative 'audiences'
+
+module Streamdal
+  class TestObj
+    include Schemas
+    include Audiences
+
+    @schemas
+    @stub
+
+    attr_accessor :stub
+
+    def initialize
+      @schemas = {}
+      @stub = RSpec::Mocks::Double.new("stub", { send_schema: nil })
+    end
+
+    def _metadata
+      {}
+    end
+  end
+end
+
+RSpec.describe "Streamdal::Schema" do
+  before(:each) do
+    @test_obj = Streamdal::TestObj.new
+
+    public_aud = Streamdal::Audience.new(1, "consume", "kafka")
+    @aud = public_aud.to_proto("test-svc")
+    expect(@test_obj.aud_to_str(@aud)).to eq("test-svc.kafka.OPERATION_TYPE_CONSUMER.consume")
+  end
+
+  it "should set and get a schema" do
+    @test_obj._set_schema(@aud, "{}")
+
+    got_schema = @test_obj._get_schema(@aud)
+    expect(got_schema).to eq("{}")
+  end
+  it "should handle schema" do
+    wasm_resp = Streamdal::Protos::WASMResponse.new
+    wasm_resp.exit_code = :WASM_EXIT_CODE_TRUE
+    wasm_resp.output_step = "{}"
+
+    @test_obj._handle_schema(@aud, Streamdal::Protos::PipelineStep.new(infer_schema: Streamdal::Protos::InferSchemaStep.new), wasm_resp)
+
+    got_schema = @test_obj._get_schema(@aud)
+    expect(got_schema).to eq("{}")
+    # sleep(1)
+    expect(@test_obj.stub).to have_received(:send_schema).with(Streamdal::Protos::SendSchemaRequest.new(audience: @aud, schema: Streamdal::Protos::Schema.new(json_schema: "{}")), metadata: @test_obj._metadata)
+  end
+end

data/lib/spec_helper.rb

@@ -0,0 +1,2 @@
+require 'simplecov'
+SimpleCov.start