code.pfad.fr/gohmekit@v0.2.1/hapip/notification/manager.go

package notification

import (
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/go-kit/log"
)

// Manager manages notifications.
// m.ConnContext & m.ConnState must be registered on the http.Server respective fields.
type Manager[Key comparable] struct {
	// If a connection is active, events will not be sent (to prevent corruption).
	// EventBuffer indicates the number of events to buffer (per connection),
	// which will be delivered when the connection is back to idle.
	EventBuffer int

	// CoalesceDuration will prevent too many events from being delivered after another.
	// Standard recommands 1s
	CoalesceDuration time.Duration
	Logger           log.Logger

	muConn sync.Mutex
	conn   map[net.Conn]*notifier

	muSubscribers sync.Mutex
	subscribers   map[Key]map[net.Conn]bool
}

func (m *Manager[Key]) HookConnEvents(server *http.Server) {
	if previous := server.ConnState; previous != nil {
		server.ConnState = func(c net.Conn, cs http.ConnState) {
			previous(c, cs)
			m.ConnState(c, cs)
		}
	} else {
		server.ConnState = m.ConnState
	}

	if previous := server.ConnContext; previous != nil {
		server.ConnContext = func(ctx context.Context, c net.Conn) context.Context {
			ctx = previous(ctx, c)
			return m.ConnContext(ctx, c)
		}
	} else {
		server.ConnContext = m.ConnContext
	}
}

type contextKey string

var contextKeyConn = contextKey("conn")

func (m *Manager[Key]) ConnContext(ctx context.Context, c net.Conn) context.Context {
	nf := &notifier{
		conn:             c,
		stateCh:          make(chan http.ConnState), // not buffered, to be blocking on ConnState
		eventCh:          make(chan io.WriterTo, m.EventBuffer),
		coalesceDuration: m.CoalesceDuration,
		logger:           m.Logger,
	}

	m.muConn.Lock()
	if m.conn == nil {
		m.conn = map[net.Conn]*notifier{c: nf}
	} else {
		m.conn[c] = nf
	}
	m.muConn.Unlock()

	go func() {
		defer m.removeFromSubscribers(c)
		defer m.removeFromConns(c)
		nf.run(c)
	}()

	return context.WithValue(ctx, contextKeyConn, c)
}

func getConnValue(ctx context.Context) net.Conn {
	conn, _ := ctx.Value(contextKeyConn).(net.Conn)
	return conn
}

func (m *Manager[Key]) ConnState(c net.Conn, s http.ConnState) {
	if s == http.StateNew {
		return
	}
	m.muConn.Lock()
	nf := m.conn[c]
	m.muConn.Unlock()

	nf.stateCh <- s
}

var errNoConn = errors.New("net.Conn not found in ctx")

func (m *Manager[Key]) Subscribe(ctx context.Context, id Key, v bool) error {
	c := getConnValue(ctx)
	if c == nil {
		return errNoConn
	}

	m.muSubscribers.Lock()
	defer m.muSubscribers.Unlock()

	switch {
	case m.subscribers == nil:
		m.subscribers = map[Key]map[net.Conn]bool{
			id: {c: v},
		}
	case m.subscribers[id] == nil:
		m.subscribers[id] = map[net.Conn]bool{c: v}
	default:
		m.subscribers[id][c] = v
	}

	return nil
}
func (m *Manager[Key]) IsSubscribed(ctx context.Context, id Key) (bool, error) {
	c := getConnValue(ctx)
	if c == nil {
		return false, errNoConn
	}

	m.muSubscribers.Lock()
	defer m.muSubscribers.Unlock()

	if m.subscribers == nil || m.subscribers[id] == nil {
		return false, nil
	}

	return m.subscribers[id][c], nil
}

// Coalescer allows event to be merged together to prevent
// sending too many events in a short time.
type Coalescer interface {
	io.WriterTo

	// Coalescer must return itself on error (not nil!)
	Coalesce(io.WriterTo) (Coalescer, error)
}

// Publish an event. If the event is a Coalescer it might be throtteled.
func (m *Manager[Key]) Publish(ctx context.Context, id Key, event io.WriterTo) {
	publisher := getConnValue(ctx)
	conns := m.subscribedConns(id, publisher)
	if len(conns) == 0 {
		return
	}

	notifiers := m.notifiers(conns)

	for _, nf := range notifiers {
		select {
		case nf.eventCh <- event:
		default:
		}
	}
}

func (m *Manager[Key]) notifiers(conns []net.Conn) []*notifier {
	notifiers := make([]*notifier, 0, len(conns))

	m.muConn.Lock()
	defer m.muConn.Unlock()

	for _, c := range conns {
		nf := m.conn[c]
		if nf == nil {
			continue
		}
		notifiers = append(notifiers, nf)
	}
	return notifiers
}

func (m *Manager[Key]) subscribedConns(id Key, publisher net.Conn) []net.Conn {
	m.muSubscribers.Lock()
	defer m.muSubscribers.Unlock()

	if m.subscribers == nil || m.subscribers[id] == nil {
		return nil
	}

	var conns []net.Conn //nolint:prealloc
	for c, v := range m.subscribers[id] {
		if !v || c == publisher {
			continue
		}
		conns = append(conns, c)
	}
	return conns
}

type notifier struct {
	conn             net.Conn
	stateCh          chan http.ConnState
	eventCh          chan io.WriterTo
	coalesceDuration time.Duration
	logger           log.Logger
}

func (n *notifier) run(c net.Conn) {
	state := http.StateNew
	var lastEvent time.Time
	for {
		// wait for the connection to be idle...
		for state != http.StateIdle {
			if state == http.StateClosed || state == http.StateHijacked {
				return
			}
			state = <-n.stateCh
		}

		select {
		case state = <-n.stateCh:
		case event := <-n.eventCh:

			nextEventDelay := n.coalesceDuration - time.Since(lastEvent)
			if nextEventDelay > 0 {
				delayExpired := time.After(nextEventDelay)
				acc, canWait := event.(Coalescer)
				for canWait {
					canWait = false
					n.logger.Log("action", "coalesce", "duration", nextEventDelay)
					select {
					case state = <-n.stateCh: // send event right away and hope for the best...
					case <-delayExpired:
					case ev := <-n.eventCh:
						var err error
						acc, err = acc.Coalesce(ev)
						if err != nil {
							n.logger.Log("action", "coalesce", "err", err)
							_, err := ev.WriteTo(c)
							if err != nil {
								n.logger.Log("action", "coalesce-fallback", "err", err)
							}
							delayExpired = time.After(n.coalesceDuration)
						}
						_, canWait = ev.(Coalescer)
					}
					event = acc
				}
			}

			_, err := event.WriteTo(c)
			lastEvent = time.Now()
			n.logger.Log("action", "write", "err", err)
		}
	}
}

func (m *Manager[Key]) removeFromSubscribers(c net.Conn) {
	m.muSubscribers.Lock()
	defer m.muSubscribers.Unlock()

	if m.subscribers == nil {
		return
	}

	for _, conns := range m.subscribers {
		delete(conns, c)
	}
}

func (m *Manager[Key]) removeFromConns(c net.Conn) {
	m.muConn.Lock()
	m.Logger.Log("action", "connection removal")
	delete(m.conn, c)
	m.muConn.Unlock()
}