llmgw/internal/cli/standbyroutes.go

package cli

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"golang.org/x/net/http2"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router"
)

// How often a standby answering requests itself hears from the leader, and
// how long it keeps routing to the replicas it last heard of when it does
// not. Long enough to ride out the leader changing hands — the engines are
// still where they were, and a dead one opens its breaker — short enough
// that a standby cut off from the cluster stops sending to a list it cannot
// check.
const (
	standbyPoll  = 2 * time.Second
	standbyTrust = time.Minute
)

// The router is what counts against limits, and shares the counts.
var _ control.RateSharer = (*router.Router)(nil)

// standbyRoutes is a standby's copy of the leader's Ready replicas, when
// standbys answer requests themselves (router.active_active). Each tick it
// reports what its router used and is answered with the replicas, in one
// request over the agent port with its control-plane certificate.
type standbyRoutes struct {
	fwd     *leaderForward
	usage   control.UsageSource
	refresh func(context.Context) error // reads the catalog again from the store
	publish func(context.Context)       // hands the router the new list now
	log     *slog.Logger
	now     func() time.Time

	mu      sync.Mutex
	reps    []control.RoutedReplica
	heard   time.Time
	catalog string               // the leader's catalog revision last seen
	sent    map[string]time.Time // the last use already reported, by replica
	want    map[string]bool      // on-demand models asked for since the last report
	kick    chan struct{}

	// The connection to the leader, kept between ticks: one per leader
	// address, the previous one's closed when the lease moves.
	client     *http.Client
	clientHost string
}

func newStandbyRoutes(fwd *leaderForward, usage control.UsageSource, refresh func(context.Context) error,
	publish func(context.Context), log *slog.Logger,
) *standbyRoutes {
	return &standbyRoutes{fwd: fwd, usage: usage, refresh: refresh, publish: publish, log: log, now: time.Now,
		sent: map[string]time.Time{}, want: map[string]bool{}, kick: make(chan struct{}, 1)}
}

// Ready satisfies control.ReplicaFeed.
func (s *standbyRoutes) Ready() ([]control.RoutedReplica, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.heard.IsZero() || s.now().Sub(s.heard) > standbyTrust {
		return nil, false
	}
	return s.reps, true
}

// serving says whether this standby has replicas to send a request to, which
// is when it keeps the client API rather than passing it to the leader.
func (s *standbyRoutes) serving() bool {
	reps, ok := s.Ready()
	return ok && len(reps) > 0
}

// wanted queues an on-demand model for the leader, and reports at once
// rather than on the next tick: someone is waiting on it.
func (s *standbyRoutes) wanted(model string) {
	s.mu.Lock()
	s.want[model] = true
	s.mu.Unlock()
	select {
	case s.kick <- struct{}{}:
	default:
	}
}

func (s *standbyRoutes) run(ctx context.Context) {
	t := time.NewTicker(standbyPoll)
	defer t.Stop()
	for {
		if s.fwd.lead.IsLeader() {
			// Its own reconciler's replicas now. What it heard as a standby
			// is forgotten, so becoming one again starts afresh.
			s.mu.Lock()
			s.reps, s.heard, s.catalog = nil, time.Time{}, ""
			clear(s.sent)
			clear(s.want)
			s.mu.Unlock()
		} else if err := s.exchange(ctx); err != nil {
			s.log.Debug("could not hear from the leader", "err", err)
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		case <-s.kick:
		}
	}
}

// exchange reports activity and takes the replicas in return.
func (s *standbyRoutes) exchange(ctx context.Context) error {
	act, sent := s.activity()
	body, err := json.Marshal(act)
	if err != nil {
		return err
	}
	target, name, err := s.fwd.leader(ctx)
	if err != nil {
		return err
	}
	client, err := s.clientFor(target.Host, target.Hostname())
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		"https://"+target.Host+forwardPrefix+"/api/v1/routing/activity", bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set(hdrForwardedBy, s.fwd.self)
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return fmt.Errorf("%s answered %s: %s", name, resp.Status, bytes.TrimSpace(b))
	}
	var rt control.Routing
	if err := json.NewDecoder(resp.Body).Decode(&rt); err != nil {
		return fmt.Errorf("%s's routing: %w", name, err)
	}

	s.mu.Lock()
	s.reps, s.heard = rt.Replicas, s.now()
	stale := rt.Catalog != s.catalog
	s.catalog = rt.Catalog
	for id, at := range sent {
		s.sent[id] = at
	}
	for _, m := range act.Want {
		delete(s.want, m)
	}
	s.mu.Unlock()

	// What the leader and the other standbys counted, so this router's
	// limits are the cluster's.
	if rs, ok := s.usage.(control.RateSharer); ok && rt.Rates != nil {
		rs.SetOtherRates(*rt.Rates)
	}

	// A model added or changed through the leader is in the store: the
	// catalog first, so the replicas published next resolve against it.
	if stale && s.refresh != nil {
		if err := s.refresh(ctx); err != nil {
			s.log.Warn("could not read the catalog the leader changed", "err", err)
			s.mu.Lock()
			s.catalog = "" // tried again next tick
			s.mu.Unlock()
		}
	}
	if s.publish != nil {
		s.publish(ctx)
	}
	return nil
}

// clientFor is the client for the leader at host, made once per leader.
// Only the run loop calls it.
func (s *standbyRoutes) clientFor(host, name string) (*http.Client, error) {
	if s.client != nil && s.clientHost == host {
		return s.client, nil
	}
	tc, err := s.fwd.tls.ClientConfig(name)
	if err != nil {
		return nil, err
	}
	if s.client != nil {
		s.client.CloseIdleConnections()
	}
	s.client = &http.Client{Transport: &http2.Transport{TLSClientConfig: tc, ReadIdleTimeout: 30 * time.Second}}
	s.clientHost = host
	return s.client, nil
}

// activity is what to report: uses newer than the last reported, and the
// models asked for. sent is what to record as reported once the leader has
// it.
func (s *standbyRoutes) activity() (control.Activity, map[string]time.Time) {
	var used map[string]time.Time
	if s.usage != nil {
		used = s.usage.LastUsed()
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	act := control.Activity{From: s.fwd.self, LastUsed: map[string]time.Time{}}
	if rs, ok := s.usage.(control.RateSharer); ok {
		r := rs.Rates()
		act.Rates = &r
	}
	for id, at := range used {
		if at.After(s.sent[id]) {
			act.LastUsed[id] = at
		}
	}
	for m := range s.want {
		act.Want = append(act.Want, m)
	}
	return act, act.LastUsed
}

// standbyActivator asks the leader for an on-demand model when a request for
// it reaches a standby: the standby's own reconciler places nothing. Which
// models load on demand, and why one cannot, are still the reconciler's to
// answer from the catalog.
type standbyActivator struct {
	*reconciler.Reconciler
	routes *standbyRoutes
}

func (a standbyActivator) Want(model string) bool {
	if !a.Reconciler.Want(model) {
		return false
	}
	if !a.routes.fwd.lead.IsLeader() {
		a.routes.wanted(model)
	}
	return true
}