Implement S2 protocol library

Co-authored-by: Bob Forma <bob@stekker.com>
Co-authored-by: Martijn Versluis <martijn@stekker.com>

Author: 3 people

README.md

@@ -118,14 +118,14 @@ json = handshake.to_json
 ```
 ┌─────────────────────┐       WebSocket       ┌─────────────────┐
 │   Your Application  │◄─────────────────────►│   CEM Server    │
-│  (Resource Manager) │                       │                 │
+│  (Resource Manager) │ ← Your custom logic   │                 │
 └──────────┬──────────┘                       └─────────────────┘
            │
            ▼
 ┌─────────────────────┐
 │ S2::Connection      │ ← Manages WebSocket lifecycle
 │ S2::Session         │ ← Handles protocol state
-│ S2::MessageHandler  │ ← Your custom logic
+│ S2::MessageHandler  │
 └─────────────────────┘
 ```
 

lib/s2.rb

@@ -0,0 +1,55 @@
+require "active_support"
+require "active_support/core_ext/class/attribute"
+require "active_support/core_ext/hash/indifferent_access"
+require "active_support/core_ext/numeric/time"
+require "active_support/core_ext/object/blank"
+require "active_support/callbacks"
+require "active_support/concern"
+require "active_support/notifications"
+require "active_support/rescuable"
+require "async"
+require "async/http/endpoint"
+require "async/notification"
+require "async/queue"
+require "async/websocket/client"
+require "dry-struct"
+require "dry-types"
+require "json"
+require "logger"
+require "securerandom"
+
+require_relative "s2/version"
+
+module S2
+  class << self
+    attr_writer :logger, :message_handler_class, :supported_protocol_versions
+
+    def logger
+      @logger ||= Logger.new(nil)
+    end
+
+    def message_handler_class
+      @message_handler_class ||= S2::MessageHandler
+    end
+
+    def supported_protocol_versions
+      @supported_protocol_versions ||= ["0.0.2-beta"]
+    end
+  end
+
+  class MessageHandler # rubocop:disable Lint/EmptyClass
+  end
+end
+
+require_relative "s2/messages"
+require_relative "s2/message_factory"
+require_relative "s2/message_sender"
+require_relative "s2/message_handler/base_error"
+require_relative "s2/message_handler/dispatching"
+require_relative "s2/message_handler/replying"
+require_relative "s2/message_handler/error_handling"
+require_relative "s2/message_handler/asserting"
+
+require_relative "s2/message_handler"
+require_relative "s2/session"
+require_relative "s2/connection"

lib/s2/connection.rb

@@ -0,0 +1,111 @@
+module S2
+  class Connection
+    INITIAL_BACKOFF = 5.seconds
+    MAX_BACKOFF = 1.hour
+
+    delegate :state,
+             to: :@session,
+             allow_nil: true
+
+    attr_reader :connected_at, :status
+
+    def initialize(resource_id:, task:, ws_url:)
+      @connected_at = nil
+      @queue = nil
+      @resource_id = resource_id
+      @session = nil
+      @task = task
+      @ws_url = ws_url
+      @status = :initialized
+      @stopping = false
+      @backoff = INITIAL_BACKOFF
+      @endpoint = Async::HTTP::Endpoint.parse(
+        ws_url,
+        alpn_protocols: Async::HTTP::Protocol::HTTP11.names,
+      )
+    end
+
+    def connect
+      until @stopping
+        begin
+          connect_and_run
+          @backoff = INITIAL_BACKOFF
+        rescue StandardError => e
+          ActiveSupport::Notifications.instrument("connection_errored.session.s2", exception: e) unless @stopping
+        ensure
+          break if @stopping
+
+          S2.logger.info "[#{self.class.name}] [#{@resource_id}] Reconnecting in #{@backoff}s..."
+          @status = :reconnecting
+          sleep @backoff
+          @backoff = [@backoff * 2, MAX_BACKOFF].min
+        end
+      end
+    end
+
+    def disconnect
+      @status = :disconnecting
+      @stopping = true
+      @session&.stop
+
+      begin
+        @queue&.close
+      rescue StandardError
+        # ignore
+      end
+
+      @status = :disconnected
+    end
+
+    def send_message(message)
+      @queue.push(message)
+    end
+
+    private
+
+    def connect_and_run
+      connect_websocket do |ws|
+        @queue = Async::Queue.new
+        @session = S2::Session.new(resource_id: @resource_id, ws:, queue: @queue)
+        @connected_at = Time.current
+        @status = :connected
+        send_handshake
+        @session.start
+      ensure
+        @session = nil
+        @queue = nil
+        @status = :disconnected
+      end
+    end
+
+    def connect_websocket(&)
+      @status = :connecting
+
+      Async::WebSocket::Client.connect(@endpoint) do |ws|
+        ActiveSupport::Notifications.instrument(
+          "connected.session.s2",
+          resource_id: @resource_id,
+          url: @ws_url,
+        )
+
+        yield ws
+      ensure
+        ActiveSupport::Notifications.instrument(
+          "disconnected.session.s2",
+          resource_id: @resource_id,
+        )
+      end
+    end
+
+    def send_handshake
+      handshake = S2::Messages::Handshake.new(
+        message_id: SecureRandom.uuid,
+        message_type: S2::Messages::HandshakeMessageType::Handshake,
+        role: S2::Messages::EnergyManagementRole::Rm,
+        supported_protocol_versions: S2.supported_protocol_versions,
+      )
+
+      send_message(handshake)
+    end
+  end
+end

lib/s2/message_factory.rb

@@ -0,0 +1,74 @@
+module S2
+  class MessageFactory
+    MESSAGE_TYPE_TO_MESSAGE_CLASS = {
+      "FRBC.ActuatorStatus" => S2::Messages::FRBCActuatorStatus,
+      "FRBC.FillLevelTargetProfile" => S2::Messages::FRBCFillLevelTargetProfile,
+      "FRBC.Instruction" => S2::Messages::FRBCInstruction,
+      "FRBC.LeakageBehaviour" => S2::Messages::FRBCLeakageBehaviour,
+      "FRBC.StorageStatus" => S2::Messages::FRBCStorageStatus,
+      "FRBC.SystemDescription" => S2::Messages::FRBCSystemDescription,
+      "FRBC.TimerStatus" => S2::Messages::FRBCTimerStatus,
+      "FRBC.UsageForecast" => S2::Messages::FRBCUsageForecast,
+      "Handshake" => S2::Messages::Handshake,
+      "HandshakeResponse" => S2::Messages::HandshakeResponse,
+      "InstructionStatusUpdate" => S2::Messages::InstructionStatusUpdate,
+      "PowerForecast" => S2::Messages::PowerForecast,
+      "PowerMeasurement" => S2::Messages::PowerMeasurement,
+      "ReceptionStatus" => S2::Messages::ReceptionStatus,
+      "ResourceManagerDetails" => S2::Messages::ResourceManagerDetails,
+      "RevokeObject" => S2::Messages::RevokeObject,
+      "SelectControlType" => S2::Messages::SelectControlType,
+      "SessionRequest" => S2::Messages::SessionRequest,
+    }.freeze
+
+    class BaseError < StandardError
+      attr_reader :message_id
+
+      def initialize(message, message_id = nil)
+        super(message)
+        @message_id = message_id
+      end
+    end
+
+    class InvalidMessageFormat < BaseError; end
+    class MissingMessageType < BaseError; end
+    class UnsupportedMessageType < BaseError; end
+    class InvalidMessagePayload < BaseError; end
+
+    class << self
+      def create_message(data)
+        message = data.is_a?(String) ? parse_json(data) : data
+        message_type = message["message_type"]
+        message_id = message["message_id"]
+
+        raise(MissingMessageType, "Message type not provided", message_id) if message_type.blank?
+
+        message_class = MESSAGE_TYPE_TO_MESSAGE_CLASS[message_type]
+
+        if message_class.nil?
+          raise(
+            UnsupportedMessageType,
+            "Message type not supported: #{message_type}",
+            message_id,
+          )
+        end
+
+        build_message(message_class, message)
+      end
+
+      private
+
+      def parse_json(data)
+        JSON.parse(data)
+      rescue JSON::ParserError
+        raise(InvalidMessageFormat, "Invalid JSON")
+      end
+
+      def build_message(message_class, message)
+        message_class.from_dynamic!(message)
+      rescue KeyError => e
+        raise(InvalidMessagePayload, "Field missing in payload: #{e.key}", message["message_id"])
+      end
+    end
+  end
+end

lib/s2/message_handler.rb

@@ -0,0 +1,71 @@
+class S2::MessageHandler
+  include ActiveSupport::Callbacks
+  include Asserting
+  include Dispatching
+  include ErrorHandling
+  include Replying
+
+  define_callbacks :handle
+  reply_with :@message_sender
+
+  STATUSES = [
+    :websocket_connected,
+    :initialized,
+  ].freeze
+
+  class << self
+    def before_handle(method_name = nil, &blk)
+      set_callback :handle, :before, method_name || wrap_callback_block(blk)
+    end
+
+    def after_handle(method_name = nil, &blk)
+      set_callback :handle, :after, method_name || wrap_callback_block(blk)
+    end
+
+    def around_handle(&)
+      set_callback :handle, :around do |_, inner|
+        instance_exec(Fiber[:s2_payload], -> { inner.call }, &)
+      end
+    end
+
+    private
+
+    def wrap_callback_block(blk)
+      raise ArgumentError, "block required" unless blk
+
+      proc do
+        instance_exec(Fiber[:s2_payload], &blk)
+      end
+    end
+  end
+
+  attr_reader :state
+
+  def initialize(message_sender:, state: { status: :websocket_connected })
+    @message_sender = message_sender
+    @state = state
+  end
+
+  def handle_message(payload)
+    Fiber[:s2_payload] = payload
+
+    run_callbacks :handle do
+      message = S2::MessageFactory.create_message(payload)
+      dispatch_message(message)
+    end
+  rescue StandardError => e
+    rescue_with_handler(e) || raise
+  end
+
+  protected
+
+  def update_state(**kwargs)
+    @state.merge!(kwargs)
+  end
+
+  def update_status(new_status)
+    raise ArgumentError, "Invalid status: #{new_status}" unless STATUSES.include?(new_status)
+
+    update_state status: new_status
+  end
+end

lib/s2/message_handler/asserting.rb

@@ -0,0 +1,46 @@
+module S2::MessageHandler::Asserting
+  extend ActiveSupport::Concern
+
+  class PermanentError < S2::MessageHandler::BaseError; end
+  class InvalidContent < S2::MessageHandler::BaseError; end
+
+  included do
+    include S2::MessageHandler::ErrorHandling
+
+    rescue_from InvalidContent do |error|
+      reply to: error.message_id, status: :invalid_content, diagnostic_label: error.message
+    end
+
+    rescue_from PermanentError do |error|
+      reply to: error.message_id, status: :permanent_error, diagnostic_label: error.message
+    end
+  end
+
+  def assert_state!(message, **expected)
+    expected.each do |key, value|
+      actual = @state[key]
+      next if actual == value
+
+      raise(
+        PermanentError.new(
+          "Invalid state: expected #{key} to be '#{value}', got '#{actual}'",
+          message.message_id,
+        ),
+      )
+    end
+  end
+
+  def assert_attribute!(message, **expected)
+    expected.each do |key, value|
+      actual = message.public_send(key)
+      next if actual == value
+
+      raise(
+        PermanentError.new(
+          "Invalid attribute: expected #{key} to be '#{value}', got '#{actual}'",
+          message.message_id,
+        ),
+      )
+    end
+  end
+end

lib/s2/message_handler/base_error.rb

@@ -0,0 +1,9 @@
+class S2::MessageHandler::BaseError < StandardError
+  attr_reader :message_id
+
+  def initialize(message, message_id)
+    super(message)
+
+    @message_id = message_id
+  end
+end