byre/internal/clients/clients.go

package clients

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/coder/websocket"

	"g.fc.run/ai/byre/internal/api"
)

const (
	readLimit    = 256 << 20
	pingInterval = 30 * time.Second
	pingTimeout  = 15 * time.Second
	expireAfter  = 24 * time.Hour
)

type Client struct {
	ID           string
	Login        string
	Host         string
	Commit       string
	ConnectedAt  time.Time
	LastActionAt *time.Time
}

type Transfer struct {
	Dir     string
	Name    string
	Receive bool
}

type conn struct {
	ws      *websocket.Conn
	done    chan struct{}
	mu      sync.Mutex
	pending map[string]chan api.ClientReply
}

type entry struct {
	Client
	conn     *conn
	seen     time.Time
	notified time.Time
}

type Registry struct {
	mu        sync.Mutex
	clients   map[string]*entry
	transfers map[string]Transfer
	changed   func()
}

func New(changed func()) *Registry {
	return &Registry{clients: map[string]*entry{}, transfers: map[string]Transfer{}, changed: changed}
}

func (r *Registry) entry(id, login, host, commit string, now time.Time) *entry {
	for k, e := range r.clients {
		if e.conn == nil && now.Sub(e.seen) > expireAfter {
			delete(r.clients, k)
		}
	}
	e := r.clients[id]
	if e == nil {
		e = &entry{Client: Client{ID: id}}
		r.clients[id] = e
	}
	e.Login = login
	e.Host = host
	e.Commit = commit
	e.seen = now
	return e
}

func (r *Registry) Acted(id, login, host, commit string) bool {
	if id == "" {
		return false
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	now := time.Now()
	e := r.entry(id, login, host, commit, now)
	e.LastActionAt = &now
	if now.Sub(e.notified) < time.Second {
		return false
	}
	e.notified = now
	return true
}

func (r *Registry) Serve(ctx context.Context, ws *websocket.Conn, id, login, host, commit string) {
	ws.SetReadLimit(readLimit)
	c := &conn{ws: ws, done: make(chan struct{}), pending: map[string]chan api.ClientReply{}}

	r.mu.Lock()
	now := time.Now().UTC()
	e := r.entry(id, login, host, commit, now)
	old := e.conn
	e.conn = c
	e.ConnectedAt = now
	r.mu.Unlock()
	if old != nil {
		old.ws.Close(websocket.StatusGoingAway, "replaced by a newer connection from this client")
	}
	slog.Info("client connected", "id", id, "host", host, "login", login, "commit", commit)
	r.changed()

	defer func() {
		r.mu.Lock()
		if e.conn == c {
			e.conn = nil
			e.seen = time.Now()
		}
		r.mu.Unlock()
		close(c.done)
		slog.Info("client disconnected", "id", id, "host", host, "login", login)
		r.changed()
	}()

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	go c.ping(ctx)

	for {
		_, b, err := ws.Read(ctx)
		if err != nil {
			return
		}
		reply := api.ClientReply{}
		err = json.Unmarshal(b, &reply)
		if err != nil {
			ws.Close(websocket.StatusUnsupportedData, "messages must be ClientReply")
			return
		}
		c.mu.Lock()
		ch := c.pending[reply.Id]
		delete(c.pending, reply.Id)
		c.mu.Unlock()
		if ch == nil {
			slog.Warn("client reply for no request", "host", host, "id", reply.Id)
			continue
		}
		ch <- reply
	}
}

func (c *conn) ping(ctx context.Context) {
	t := time.NewTicker(pingInterval)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			pctx, cancel := context.WithTimeout(ctx, pingTimeout)
			err := c.ws.Ping(pctx)
			cancel()
			if err != nil {
				c.ws.Close(websocket.StatusGoingAway, "ping timed out")
				return
			}
		}
	}
}

func (r *Registry) List() []Client {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := []Client{}
	for _, e := range r.clients {
		if e.conn != nil {
			out = append(out, e.Client)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ConnectedAt.Before(out[j].ConnectedAt) })
	return out
}

func (r *Registry) Call(ctx context.Context, name string, req api.ClientRequest) (api.ClientReply, error) {
	r.mu.Lock()
	matches := []*conn{}
	ids := []string{}
	connected := []string{}
	for _, e := range r.clients {
		if e.conn == nil {
			continue
		}
		connected = append(connected, fmt.Sprintf("%s (%s)", e.Host, e.ID))
		if e.ID == name || e.Host == name {
			matches = append(matches, e.conn)
			ids = append(ids, e.ID)
		}
	}
	r.mu.Unlock()
	slices.Sort(connected)
	if len(matches) == 0 {
		return api.ClientReply{}, fmt.Errorf("no client named %q is connected; connected: %s", name, strings.Join(connected, ", "))
	}
	if len(matches) > 1 {
		slices.Sort(ids)
		return api.ClientReply{}, fmt.Errorf("%d clients on %s are connected; name one by its id: %s", len(matches), name, strings.Join(ids, ", "))
	}
	c := matches[0]

	req.Id = rand.Text()
	ch := make(chan api.ClientReply, 1)
	c.mu.Lock()
	c.pending[req.Id] = ch
	c.mu.Unlock()
	defer func() {
		c.mu.Lock()
		delete(c.pending, req.Id)
		c.mu.Unlock()
	}()

	data, err := json.Marshal(req)
	if err != nil {
		return api.ClientReply{}, err
	}
	err = c.ws.Write(ctx, websocket.MessageText, data)
	if err != nil {
		return api.ClientReply{}, fmt.Errorf("send to %s: %w", name, err)
	}

	select {
	case reply := <-ch:
		if reply.Error != nil {
			return reply, errors.New(*reply.Error)
		}
		return reply, nil
	case <-c.done:
		return api.ClientReply{}, fmt.Errorf("%s disconnected before replying", name)
	case <-ctx.Done():
		return api.ClientReply{}, fmt.Errorf("%s didn't reply: %w", name, ctx.Err())
	}
}

func (r *Registry) AddTransfer(t Transfer) string {
	id := rand.Text()
	r.mu.Lock()
	defer r.mu.Unlock()
	r.transfers[id] = t
	return id
}

func (r *Registry) RemoveTransfer(id string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	delete(r.transfers, id)
}

func (r *Registry) Transfer(id string) (Transfer, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	t, ok := r.transfers[id]
	return t, ok
}