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/misc/byre/internal/api"
)

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

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

type Transfer struct {
	Dir  string
	Name string
}

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

type Registry struct {
	mu        sync.Mutex
	conns     map[string]*conn
	transfers map[string]Transfer
}

func New() *Registry {
	return &Registry{conns: map[string]*conn{}, transfers: map[string]Transfer{}}
}

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

	r.mu.Lock()
	old := r.conns[client.Host]
	r.conns[client.Host] = c
	r.mu.Unlock()
	if old != nil {
		old.ws.Close(websocket.StatusGoingAway, "replaced by a newer connection from this host")
	}
	slog.Info("client connected", "host", client.Host, "login", client.Login, "commit", client.Commit)

	defer func() {
		r.mu.Lock()
		if r.conns[client.Host] == c {
			delete(r.conns, client.Host)
		}
		r.mu.Unlock()
		close(c.done)
		slog.Info("client disconnected", "host", client.Host, "login", client.Login)
	}()

	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", client.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 _, c := range r.conns {
		out = append(out, c.Client)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Host < out[j].Host })
	return out
}

func (r *Registry) Call(ctx context.Context, host string, req api.ClientRequest) (api.ClientReply, error) {
	r.mu.Lock()
	c := r.conns[host]
	hosts := []string{}
	for h := range r.conns {
		hosts = append(hosts, h)
	}
	r.mu.Unlock()
	if c == nil {
		slices.Sort(hosts)
		return api.ClientReply{}, fmt.Errorf("no client named %q is connected; connected: %s", host, strings.Join(hosts, ", "))
	}

	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", host, 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", host)
	case <-ctx.Done():
		return api.ClientReply{}, fmt.Errorf("%s didn't reply: %w", host, 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
}