pull websockets out of engine

engine.go

@@ -4,7 +4,6 @@ import (
 	"log/slog"
 	"net/http"
 
-	"github.com/gorilla/websocket"
 	"github.com/wisplite/tether/reactivity"
 	"gorm.io/gorm"
 )
@@ -13,11 +12,13 @@ type Engine struct {
 	db        *gorm.DB
 	mutations map[string]func(ctx *MutationCtx) error
 	queries   map[string]func(ctx *QueryCtx) error
+	tracker   *reactivity.Tracker
 }
 
 func NewEngine(db *gorm.DB) *Engine {
 	slog.SetLogLoggerLevel(slog.LevelDebug)
-	return &Engine{db: db, mutations: make(map[string]func(ctx *MutationCtx) error), queries: make(map[string]func(ctx *QueryCtx) error)}
+	tracker := reactivity.NewTracker()
+	return &Engine{db: db, mutations: make(map[string]func(ctx *MutationCtx) error), queries: make(map[string]func(ctx *QueryCtx) error), tracker: tracker}
 }
 
 func (e *Engine) RegisterMutation(name string, mutation func(ctx *MutationCtx) error) {
@@ -36,22 +37,22 @@ func (e *Engine) CreateTable(name string, schema interface{}) {
 }
 
 func (e *Engine) Handle(w http.ResponseWriter, r *http.Request) {
-	reactivity.Handle(w, r, e) // wraps the raw websocket connection with the engine handler
+	reactivity.Handle(w, r, e, e.tracker) // wraps the raw websocket connection with the engine handler
 }
 
-func (e *Engine) OnConnect(conn *websocket.Conn) error {
+func (e *Engine) OnConnect(clientID string) error {
 	slog.Debug("Connected to websocket")
 	// TODO: implement the logic to handle the connection
 	return nil
 }
 
-func (e *Engine) OnDisconnect(conn *websocket.Conn) error {
+func (e *Engine) OnDisconnect(clientID string) error {
 	slog.Debug("Disconnected from websocket")
 	// TODO: implement the logic to handle the disconnection
 	return nil
 }
 
-func (e *Engine) OnReceiveMessage(msg map[string]interface{}) error {
+func (e *Engine) OnReceiveMessage(clientID string, msg map[string]interface{}) error {
 	slog.Debug("Received message", "message", msg)
 	// TODO: implement the logic to handle the message
 	return nil

reactivity/websocket.go

@@ -10,14 +10,14 @@ import (
 
 // MessageReceiver receives decoded WebSocket payloads. Implemented by tether.Engine.
 type EngineHandler interface {
-	OnConnect(conn *websocket.Conn) error
-	OnDisconnect(conn *websocket.Conn) error
-	OnReceiveMessage(msg map[string]interface{}) error
+	OnConnect(clientID string) error
+	OnDisconnect(clientID string) error
+	OnReceiveMessage(clientID string, msg map[string]interface{}) error
 }
 
 var upgrader = websocket.Upgrader{}
 
-func Handle(w http.ResponseWriter, r *http.Request, e EngineHandler) {
+func Handle(w http.ResponseWriter, r *http.Request, e EngineHandler, tracker *Tracker) {
 	ws, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		slog.Error("WS: Failed to upgrade to websocket", "error", err)
@@ -25,13 +25,14 @@ func Handle(w http.ResponseWriter, r *http.Request, e EngineHandler) {
 	}
 	defer ws.Close()
 
-	err = e.OnConnect(ws)
+	client := NewClient(ws)
+	err = e.OnConnect(client.ID)
 	if err != nil {
 		slog.Error("WS: Failed to call onConnect handler", "error", err)
 		return
 	}
-	defer e.OnDisconnect(ws)
-
+	defer e.OnDisconnect(client.ID)
+	tracker.Track(client)
 	for {
 		_, message, err := ws.ReadMessage()
 		if err != nil {
@@ -47,7 +48,7 @@ func Handle(w http.ResponseWriter, r *http.Request, e EngineHandler) {
 			return
 		}
 		slog.Debug("WS: Unmarshalled message", "message", msg)
-		err = e.OnReceiveMessage(msg)
+		err = e.OnReceiveMessage(client.ID, msg)
 		if err != nil {
 			slog.Error("WS: Failed to on receive message", "error", err)
 			return