htmxmqtt/handler.go

package htmxmqtt

import (
  "time"
  "net/http"
  "log/slog"
  "fmt"
  "context"
  "sync"
  "slices"

  "nhooyr.io/websocket"
  mqtt "github.com/eclipse/paho.mqtt.golang"
)

type MQTTFormatFunc func(mqtt.Message) []byte
type MQTTHandler struct {
  sync.Mutex
  format MQTTFormatFunc
  channels []chan mqtt.Message
}

func (handler *MQTTHandler) processMessage(client mqtt.Client, message mqtt.Message) {
  message.Ack()
  handler.Lock()
  defer handler.Unlock()

  remaining := make([]chan mqtt.Message, 0, len(handler.channels))
  for _, channel := range(handler.channels) {
    select {
    case channel <- message:
      remaining = append(remaining, channel)
    default:
      slog.Warn("Channel overflow")
    }
  }

  handler.channels = remaining
}

func (handler *MQTTHandler) addChannel(channel chan mqtt.Message) func() {
  handler.Lock()
  defer handler.Unlock()

  handler.channels = append(handler.channels, channel)
  return func() {
    handler.Lock()
    defer handler.Unlock()

    idx := slices.Index(handler.channels, channel)
    if idx < 0 {
      return
    }

    handler.channels[idx] = handler.channels[len(handler.channels)-1]
    handler.channels = handler.channels[:len(handler.channels)-1]
  }
}

type MQTTHandlerClient struct {
  mqtt.Client
  subscriptions map[*MQTTHandler]string
  subscribeTimeout time.Duration
}

func NewMQTTHandlerClient(broker string, username string, password string, id string) (*MQTTHandlerClient, error) {
  opts := mqtt.NewClientOptions()
  opts.AddBroker(broker)
  opts.SetClientID(id)
  opts.SetUsername(username)
  opts.SetPassword(password)
  opts.SetKeepAlive(2 * time.Second)
  opts.SetPingTimeout(1 * time.Second)
  client := mqtt.NewClient(opts)

  if token := client.Connect(); token.Wait() && token.Error() != nil {
    return nil, token.Error()
  }

  return &MQTTHandlerClient{
    Client: client,
    subscriptions: map[*MQTTHandler]string{},
    subscribeTimeout: 1*time.Second,
  }, nil
}

func (client *MQTTHandlerClient) NewHandler(subscription string, format MQTTFormatFunc) (*MQTTHandler, error) {
  handler := &MQTTHandler{
    format: format,
  }

  sub_token := client.Subscribe(subscription, 0x00, handler.processMessage)

  timeout := sub_token.WaitTimeout(client.subscribeTimeout)
  if timeout == false || sub_token.Error() != nil {
    return nil, fmt.Errorf("Failed to subscribe to %s - %e", subscription, sub_token.Error())
  }

  client.subscriptions[handler] = subscription
  return handler, nil
}

func (handler *MQTTHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  channel := make(chan mqtt.Message, 1)
  remove_channel := handler.addChannel(channel)
  defer remove_channel()

  conn, err := websocket.Accept(w, r, nil)
  if err != nil {
    slog.Error("websocket accept error", "error", err)
    return
  } else {
    slog.Info("new websocket connection", "addr", r.RemoteAddr)
  }

  defer conn.CloseNow()
  ctx, cancel_func := context.WithCancel(context.Background())

  go func(conn *websocket.Conn, cancel_func context.CancelFunc) {
    for true {
      msg_type, data, err := conn.Read(ctx)
      if err != nil {
        slog.Error("websocket error", "error", err)
        cancel_func()
        break
      } else {
        slog.Debug("websocket data", "type", msg_type, "data", string(data))
      }
    }
  }(conn, cancel_func)

  running := true
  done := ctx.Done()
  for running == true {
    select {
    case <- done:
      slog.Debug("websocket context done")
      running = false
    case message := <- channel:
      text := handler.format(message)
      slog.Debug("websocket write", "data", text)
      err := conn.Write(ctx, websocket.MessageText, text)
      if err != nil {
        slog.Error("websocket write error", "error", err)
        running = false
      }
    }
  }

  slog.Info("closing websocket", "addr", r.RemoteAddr)
}

func PayloadFormatFunc(template string) MQTTFormatFunc {
  return func(message mqtt.Message) []byte {
    return []byte(fmt.Sprintf(template, message.Payload()))
  }
}