vibeproxy/internal/streams/manager.go

package streams

import (
	"context"
	"database/sql"
	"log/slog"
	"reflect"
	"sync"

	"github.com/owen/vibeproxy/internal/database"
	"github.com/owen/vibeproxy/internal/models"
)

type running struct {
	model  models.Stream
	cancel context.CancelFunc
	tcp    *TCPProxy
	udp    *UDPProxy
}

// Manager owns the set of currently-listening TCP/UDP streams and can
// reconcile it against the database without disturbing streams that
// haven't changed, so existing connections on unaffected streams are never
// dropped by a reload.
type Manager struct {
	db  *sql.DB
	log *slog.Logger

	mu      sync.Mutex
	active  map[int64]*running // stream ID -> running listener
	baseCtx context.Context
}

func NewManager(db *sql.DB, log *slog.Logger) *Manager {
	return &Manager{db: db, log: log, active: map[int64]*running{}, baseCtx: context.Background()}
}

// Reload fetches the current stream configuration and starts/stops
// listeners to match it. Streams whose configuration is unchanged are left
// running untouched.
func (m *Manager) Reload(ctx context.Context) error {
	streams, err := database.ListStreams(m.db)
	if err != nil {
		return err
	}
	desired := map[int64]models.Stream{}
	for _, s := range streams {
		if s.Enabled {
			desired[s.ID] = s
		}
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	// Stop streams that were removed or disabled.
	for id, r := range m.active {
		if _, ok := desired[id]; !ok {
			m.stopLocked(id, r)
		}
	}

	for id, s := range desired {
		if r, ok := m.active[id]; ok {
			if reflect.DeepEqual(r.model, s) {
				continue // unchanged, leave the listener (and its connections) alone
			}
			m.stopLocked(id, r)
		}
		m.startLocked(id, s)
	}
	return nil
}

func (m *Manager) startLocked(id int64, s models.Stream) {
	ctx, cancel := context.WithCancel(m.baseCtx)
	r := &running{model: s, cancel: cancel}
	var err error
	switch s.Protocol {
	case models.StreamTCP:
		r.tcp = NewTCPProxy(s, m.log)
		err = r.tcp.Start(ctx)
	case models.StreamUDP:
		r.udp = NewUDPProxy(s, m.log)
		err = r.udp.Start(ctx)
	}
	if err != nil {
		m.log.Error("failed to start stream", "id", id, "protocol", s.Protocol, "err", err)
		cancel()
		return
	}
	m.active[id] = r
}

func (m *Manager) stopLocked(id int64, r *running) {
	r.cancel()
	if r.tcp != nil {
		r.tcp.Stop()
	}
	if r.udp != nil {
		r.udp.Stop()
	}
	delete(m.active, id)
}

// StopAll shuts down every running stream (used on process shutdown).
func (m *Manager) StopAll() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for id, r := range m.active {
		m.stopLocked(id, r)
	}
}