teraWeb/backend/internal/directstream/ticket_pool.go

package directstream

import (
	"context"
	"errors"
	"strings"
	"sync"
)

type ticketEndpoint interface {
	TicketConsumer
	TicketLifecycle
	Status(context.Context) (TicketServiceSnapshot, error)
}

// TicketSocketPool lets one tera-sessiond broker accept Direct Mode tickets
// issued by more than one tera-api process. Each consumed binding is pinned to
// the authority that accepted it so heartbeats, events, writer release, and
// close requests return to the same API process.
type TicketSocketPool struct {
	endpoints []ticketEndpoint

	mu       sync.RWMutex
	bindings map[TicketBinding]ticketEndpoint
}

func NewTicketSocketPool(paths []string) *TicketSocketPool {
	endpoints := make([]ticketEndpoint, 0, len(paths))
	seen := make(map[string]struct{}, len(paths))
	for _, path := range paths {
		path = strings.TrimSpace(path)
		if path == "" {
			continue
		}
		if _, exists := seen[path]; exists {
			continue
		}
		seen[path] = struct{}{}
		endpoints = append(endpoints, NewTicketSocketClient(path))
	}
	if len(endpoints) == 0 {
		endpoints = append(endpoints, NewTicketSocketClient(DefaultTicketSocketPath))
	}
	return newTicketSocketPool(endpoints)
}

func newTicketSocketPool(endpoints []ticketEndpoint) *TicketSocketPool {
	return &TicketSocketPool{
		endpoints: append([]ticketEndpoint(nil), endpoints...),
		bindings:  make(map[TicketBinding]ticketEndpoint),
	}
}

func (p *TicketSocketPool) ConsumeTicket(ctx context.Context, ticket string, binding TicketBinding) (TicketClaims, error) {
	var failures []error
	for _, endpoint := range p.endpoints {
		claims, err := endpoint.ConsumeTicket(ctx, ticket, binding)
		if err != nil {
			failures = append(failures, err)
			continue
		}
		p.mu.Lock()
		p.bindings[binding] = endpoint
		p.mu.Unlock()
		return claims, nil
	}
	if len(failures) == 0 {
		return TicketClaims{}, ErrInvalidTicket
	}
	return TicketClaims{}, errors.Join(failures...)
}

func (p *TicketSocketPool) ReleaseDirectWriter(ctx context.Context, binding TicketBinding) error {
	endpoint, err := p.endpoint(binding)
	if err != nil {
		return err
	}
	return endpoint.ReleaseDirectWriter(ctx, binding)
}

func (p *TicketSocketPool) HeartbeatDirectStream(ctx context.Context, binding TicketBinding) error {
	endpoint, err := p.endpoint(binding)
	if err != nil {
		return err
	}
	return endpoint.HeartbeatDirectStream(ctx, binding)
}

func (p *TicketSocketPool) RecordDirectStreamEvent(ctx context.Context, binding TicketBinding, event SessionEvent) error {
	endpoint, err := p.endpoint(binding)
	if err != nil {
		return err
	}
	return endpoint.RecordDirectStreamEvent(ctx, binding, event)
}

func (p *TicketSocketPool) CloseDirectStream(ctx context.Context, binding TicketBinding) error {
	p.mu.Lock()
	endpoint, exists := p.bindings[binding]
	if exists {
		delete(p.bindings, binding)
	}
	p.mu.Unlock()
	if !exists {
		return ErrTicketBindingInactive
	}
	return endpoint.CloseDirectStream(ctx, binding)
}

func (p *TicketSocketPool) Status(ctx context.Context) (TicketServiceSnapshot, error) {
	var result TicketServiceSnapshot
	successes := 0
	var failures []error
	for _, endpoint := range p.endpoints {
		status, err := endpoint.Status(ctx)
		if err != nil {
			failures = append(failures, err)
			continue
		}
		successes++
		result.ActiveBindings += status.ActiveBindings
		result.MaxActiveBindings += status.MaxActiveBindings
		result.Consumed += status.Consumed
		result.Heartbeats += status.Heartbeats
		result.ReleasedWriters += status.ReleasedWriters
		result.ClosedStreams += status.ClosedStreams
		result.ExpiredStreams += status.ExpiredStreams
		result.LifecycleFailures += status.LifecycleFailures
		result.SessionEvents += status.SessionEvents
		result.SessionEventFailures += status.SessionEventFailures
	}
	if successes > 0 {
		return result, nil
	}
	if len(failures) == 0 {
		return result, errors.New("no direct stream ticket authorities are configured")
	}
	return result, errors.Join(failures...)
}

func (p *TicketSocketPool) endpoint(binding TicketBinding) (ticketEndpoint, error) {
	p.mu.RLock()
	endpoint, exists := p.bindings[binding]
	p.mu.RUnlock()
	if !exists {
		return nil, ErrTicketBindingInactive
	}
	return endpoint, nil
}