Puma WebSocket handler

websocket.rb

class Wrapper
  include Logging

  def initialize(driver, socket)
    @driver = driver
    @socket = socket
    @queue = Queue.new
    @closed = false

    setup_driver
    start_driver_thread
  end

  def read
    loop do
      if @closed && @queue.empty?
        raise EOFError, "WebSocket closed"
      end

      # Attempt to pop from the queue
      message = @queue.pop

      if message
        logger.debug "Wrapper: Dequeued message #{message.inspect}"
        return message
      end
    end
  end

  def write(data)
    logger.debug "Wrapper: Writing data #{data.inspect}"
    @driver.binary(data)
  end

  def close
    logger.debug "Wrapper: Closing"
    @closed = true
    @reader_thread&.join if Thread.current != @reader_thread # Ensure the thread is properly terminated if called from outside
    @driver.close
    @socket.close unless @socket.closed?
  end

  private

  def setup_driver
    @driver.on(:message) { |event| @queue << event.data }
    @driver.on(:close) do
      @closed = true
    end

    logger.debug "Wrapper: Starting WebSocket driver"
    @driver.start
  end

  def start_driver_thread
    @reader_thread = Thread.new do
      logger.debug "Wrapper: Starting reader"
      while !@closed
        logger.debug "Wrapper: Reading data"
        data = @socket.readpartial(1024)
        @driver.parse(data)
      end
    rescue EOFError, IOError
      close
    ensure
      @closed = true
    end
  end
end

class Client
  include Logging
  attr_reader :url, :websocket

  def initialize(url)
    @url = url
    @uri = URI.parse(url)

    start
  end

  def start
    logger.debug "Client: Connecting to #{@url}"
    @socket = TCPSocket.new(@uri.host, @uri.port)
    driver = ::WebSocket::Driver.client(self)
    @websocket = Terminalwire::WebSocket::Wrapper.new(driver, @socket)
    logger.debug "Client: Connection established"
  rescue => e
    logger.error "Client: Failed to connect to WebSocket server: #{e.message}"
    raise
  end

  def close
    logger.debug "Client: Stopping"
    @websocket.close
    @socket.close
  end

  def write(data)
    logger.debug "Client Writing: #{data}"
    @socket.write(data)
  end
end