visionA/local-agent/server/internal/camera/mjpeg.go

package camera

import (
	"fmt"
	"net/http"
	"sync"
)

type MJPEGStreamer struct {
	frameCh chan []byte
	clients map[chan []byte]bool
	mu      sync.Mutex
}

func NewMJPEGStreamer() *MJPEGStreamer {
	return &MJPEGStreamer{
		frameCh: make(chan []byte, 3),
		clients: make(map[chan []byte]bool),
	}
}

func (s *MJPEGStreamer) FrameChannel() chan<- []byte {
	return s.frameCh
}

func (s *MJPEGStreamer) Run() {
	for frame := range s.frameCh {
		s.mu.Lock()
		for ch := range s.clients {
			select {
			case ch <- frame:
			default:
				// drop frame for slow client
			}
		}
		s.mu.Unlock()
	}
}

func (s *MJPEGStreamer) AddClient() chan []byte {
	ch := make(chan []byte, 3)
	s.mu.Lock()
	s.clients[ch] = true
	s.mu.Unlock()
	return ch
}

func (s *MJPEGStreamer) RemoveClient(ch chan []byte) {
	s.mu.Lock()
	delete(s.clients, ch)
	s.mu.Unlock()
}

func (s *MJPEGStreamer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "multipart/x-mixed-replace; boundary=frame")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming not supported", 500)
		return
	}

	clientCh := s.AddClient()
	defer s.RemoveClient(clientCh)

	for {
		select {
		case <-r.Context().Done():
			return
		case frame := <-clientCh:
			fmt.Fprintf(w, "--frame\r\n")
			fmt.Fprintf(w, "Content-Type: image/jpeg\r\n")
			fmt.Fprintf(w, "Content-Length: %d\r\n\r\n", len(frame))
			w.Write(frame)
			fmt.Fprintf(w, "\r\n")
			flusher.Flush()
		}
	}
}