go-otel/internal/websocket/server.go

// Package websocket handles WebSocket connections and broadcasting telemetry updates to clients.
package websocket

import (
	"encoding/json"
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

// DataType represents the type of telemetry data.
type DataType string

// Telemetry data types.
const (
	DataTypeTraces  DataType = "traces"
	DataTypeMetrics DataType = "metrics"
	DataTypeLogs    DataType = "logs"
)

// HistoryProvider provides access to historical telemetry entries.
type HistoryProvider interface {
	GetHistory(limit int) []HistoryEntry
}

// HistoryEntry represents a single archived telemetry entry.
type HistoryEntry struct {
	Type      string          `json:"type"`
	Timestamp int64           `json:"timestamp"`
	Data      json.RawMessage `json:"data"`
}

// Entry represents a telemetry update sent to clients.
type Entry struct {
	Type string          `json:"type"`
	Data json.RawMessage `json:"data"`
}

// Broadcaster interface for sending telemetry updates to all connected clients.
type Broadcaster interface {
	Broadcast(data *Entry)
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// Server manages WebSocket client connections and broadcasts.
type Server struct {
	clients         map[*websocket.Conn]bool
	mutex           sync.Mutex
	historyProvider HistoryProvider
}

// New creates a new WebSocket Server instance.
func New(historyProvider HistoryProvider) *Server {
	return &Server{

		clients:         make(map[*websocket.Conn]bool),
		historyProvider: historyProvider,
	}
}

// HandleConnections upgrades HTTP requests to WebSocket and manages client lifecycle.
func (s *Server) HandleConnections(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("failed to upgrade connection: %v", err)
		return
	}
	defer func() {
		if err := ws.Close(); err != nil {
			log.Printf("failed to close connection: %v", err)
		}
	}()

	s.mutex.Lock()
	s.clients[ws] = true
	s.mutex.Unlock()

	log.Printf("client connected")

	// // Send history to new client
	// if s.historyProvider != nil {
	// 	history := s.historyProvider.GetHistory(50)
	// 	for _, entry := range history {
	// 		if data, err := json.Marshal(entry); err == nil {
	// 			if err := ws.WriteMessage(websocket.TextMessage, data); err != nil {
	// 				log.Printf("failed to send history to client: %v", err)
	// 				break
	// 			}
	// 		}
	// 	}
	// }

	for {
		// Read message from browser
		_, _, err := ws.ReadMessage()
		if err != nil {
			log.Printf("read error: %v", err)
			break
		}
	}

	s.mutex.Lock()
	delete(s.clients, ws)
	s.mutex.Unlock()
	log.Printf("client disconnected")
}

// Broadcast sends a telemetry entry to all connected clients.
func (s *Server) Broadcast(entry *Entry) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	msg, err := json.Marshal(entry)
	if err != nil {
		log.Printf("failed to marshal log entry: %v", err)
	} else {
		for client := range s.clients {
			if err := client.WriteMessage(websocket.TextMessage, msg); err != nil {
				log.Printf("write error: %v", err)
			}
		}
	}
}

// SendHistory sends historical entries to a specific WebSocket client.
func (s *Server) SendHistory(ws *websocket.Conn) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	// Default to 100 entries for new socket connections if not otherwise specified
	history := s.historyProvider.GetHistory(100)
	for _, entry := range history {
		data, err := json.Marshal(entry)
		if err != nil {
			log.Printf("failed to marshal history entry: %v", err)
			continue
		}
		if err := ws.WriteMessage(websocket.TextMessage, data); err != nil {
			log.Printf("write error: %v", err)
		}
	}
}

// HandleMessage processes incoming messages from a WebSocket client.
func (s *Server) HandleMessage(ws *websocket.Conn, msg []byte) {
	var message map[string]interface{}
	if err := json.Unmarshal(msg, &message); err != nil {
		log.Printf("error unmarshaling message: %v", err)
		return
	}

	// Handle the message (this is just an example, adjust according to your needs)
	switch message["type"] {
	case "ping":
		log.Printf("received ping")
		s.SendPong(ws)
	default:
		log.Printf("unknown message type : %v", message)
	}
}

// SendPong sends a pong message back to a client.
func (s *Server) SendPong(ws *websocket.Conn) {
	if err := ws.WriteMessage(websocket.PongMessage, nil); err != nil {
		log.Printf("error sending pong: %v", err)
	}
}