275 lines
8.8 KiB
Go
275 lines
8.8 KiB
Go
package hub
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
"github.com/coder/websocket/wsjson"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/shcizo/package-updater/internal/master/events"
|
|
"github.com/shcizo/package-updater/internal/master/model"
|
|
"github.com/shcizo/package-updater/internal/protocol"
|
|
)
|
|
|
|
var t0 = time.Date(2026, 10, 12, 12, 0, 0, 0, time.UTC)
|
|
|
|
func sendJob(t *testing.T, c *websocket.Conn, typ, id string, payload any) {
|
|
t.Helper()
|
|
env, err := protocol.EncodeJob(typ, id, payload)
|
|
require.NoError(t, err)
|
|
require.NoError(t, wsjson.Write(context.Background(), c, env))
|
|
}
|
|
|
|
func readEnv(t *testing.T, c *websocket.Conn) protocol.Envelope {
|
|
t.Helper()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
var env protocol.Envelope
|
|
require.NoError(t, wsjson.Read(ctx, c, &env))
|
|
return env
|
|
}
|
|
|
|
func waitRoute(t *testing.T, h *Hub) model.Route {
|
|
t.Helper()
|
|
svc := model.Service{Kind: protocol.KindCompose, AgentID: ptr(node2.ID)}
|
|
var r model.Route
|
|
require.Eventually(t, func() bool {
|
|
var err error
|
|
r, err = h.Route(svc)
|
|
return err == nil
|
|
}, 2*time.Second, 5*time.Millisecond)
|
|
return r
|
|
}
|
|
|
|
func TestHub_ReconcilesBeforeRoutable(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
store.lostGate = make(chan struct{})
|
|
store.lostResult = []int64{5}
|
|
h, pub, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2", ProtocolVersion: 2, KnownJobs: []string{"7", "9", "x"}})
|
|
svc := model.Service{Kind: protocol.KindCompose, AgentID: ptr(node2.ID)}
|
|
time.Sleep(50 * time.Millisecond) // hello is now blocked inside FailLostJobs
|
|
_, err = h.Route(svc)
|
|
assert.ErrorIs(t, err, model.ErrAgentOffline)
|
|
|
|
close(store.lostGate)
|
|
waitRoute(t, h)
|
|
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
assert.Equal(t, [][]int64{{7, 9}}, store.lostCalls)
|
|
assert.True(t, pub.has(events.Jobs))
|
|
assert.True(t, pub.has(events.JobTopic(5)))
|
|
}
|
|
|
|
func TestHub_StartJobSendsEnvelope(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
h, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
r := waitRoute(t, h)
|
|
|
|
start := protocol.JobStart{Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web", Ref: "r:1"}
|
|
require.NoError(t, h.StartJob(r, 42, start))
|
|
|
|
env := readEnv(t, c)
|
|
assert.Equal(t, protocol.TypeJobStart, env.Type)
|
|
assert.Equal(t, "42", env.ID)
|
|
got, err := protocol.Decode[protocol.JobStart](env)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, start, got)
|
|
}
|
|
|
|
func TestHub_StartJobOnReplacedConnectionFails(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
h, _, srv := newHub(t, store)
|
|
first, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer first.CloseNow()
|
|
send(t, first, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
old := waitRoute(t, h)
|
|
|
|
second, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer second.CloseNow()
|
|
|
|
// dial returns before the hub has registered the connection, so wait for
|
|
// the replacement to take effect: the old route is then refused. Until the
|
|
// new connection says hello, nothing is routable.
|
|
require.Eventually(t, func() bool {
|
|
return errors.Is(h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)
|
|
}, 2*time.Second, 5*time.Millisecond)
|
|
send(t, second, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
fresh := waitRoute(t, h)
|
|
require.NotEqual(t, old.ConnID, fresh.ConnID)
|
|
assert.ErrorIs(t, h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)
|
|
|
|
require.NoError(t, h.StartJob(fresh, 2, protocol.JobStart{Ref: "r:2"}))
|
|
env := readEnv(t, second)
|
|
assert.Equal(t, "2", env.ID, "the new connection got only the job sent to it")
|
|
}
|
|
|
|
func TestHub_ProgressAndResult(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
store.progressOK, store.finishOK = true, true
|
|
_, pub, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
|
|
sendJob(t, c, protocol.TypeJobProgress, "42", protocol.JobProgress{Output: "out"})
|
|
require.Eventually(t, func() bool { return pub.has(events.JobTopic(42)) }, time.Second, 5*time.Millisecond)
|
|
store.mu.Lock()
|
|
assert.Equal(t, []string{"2|42|out"}, store.progress)
|
|
store.mu.Unlock()
|
|
|
|
pub.reset()
|
|
res := protocol.JobResult{Status: protocol.ResultSucceeded, Ref: "r:1@sha256:d", Output: "done", StartedAt: t0, FinishedAt: t0}
|
|
sendJob(t, c, protocol.TypeJobResult, "42", res)
|
|
ack := readEnv(t, c)
|
|
assert.Equal(t, protocol.TypeJobAck, ack.Type)
|
|
assert.Equal(t, "42", ack.ID)
|
|
assert.True(t, pub.has(events.Services))
|
|
assert.True(t, pub.has(events.JobTopic(42)))
|
|
store.mu.Lock()
|
|
require.Len(t, store.finishes, 1)
|
|
assert.Equal(t, model.JobSucceeded, store.finishes[0].Status)
|
|
assert.Equal(t, "r:1@sha256:d", store.finishes[0].Ref)
|
|
store.mu.Unlock()
|
|
|
|
// Not applied (already recorded, or not this agent's): still acked.
|
|
store.mu.Lock()
|
|
store.finishOK = false
|
|
store.mu.Unlock()
|
|
pub.reset()
|
|
sendJob(t, c, protocol.TypeJobResult, "43", res)
|
|
assert.Equal(t, "43", readEnv(t, c).ID)
|
|
assert.False(t, pub.has(events.JobTopic(43)))
|
|
|
|
// Non-numeric ID: acked, no store call.
|
|
store.mu.Lock()
|
|
before := len(store.finishes)
|
|
store.mu.Unlock()
|
|
sendJob(t, c, protocol.TypeJobResult, "abc", res)
|
|
assert.Equal(t, "abc", readEnv(t, c).ID)
|
|
store.mu.Lock()
|
|
assert.Len(t, store.finishes, before)
|
|
store.mu.Unlock()
|
|
|
|
// Unknown status is stored as a failure.
|
|
store.mu.Lock()
|
|
store.finishOK = true
|
|
store.mu.Unlock()
|
|
sendJob(t, c, protocol.TypeJobResult, "44", protocol.JobResult{Status: "weird", Ref: "r:9"})
|
|
assert.Equal(t, "44", readEnv(t, c).ID)
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
last := store.finishes[len(store.finishes)-1]
|
|
assert.Equal(t, model.JobFailed, last.Status)
|
|
assert.Equal(t, `agent sent invalid status "weird"`, last.Error)
|
|
assert.Empty(t, last.Ref)
|
|
}
|
|
|
|
func TestHub_ResultStoreErrorClosesConnection(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
store.finishErr = assert.AnError
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
sendJob(t, c, protocol.TypeJobResult, "42", protocol.JobResult{Status: protocol.ResultFailed})
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
_, _, err = c.Read(ctx)
|
|
assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err), "no ack, connection closed")
|
|
}
|
|
|
|
func TestHub_IgnoresUnknownType(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
require.NoError(t, wsjson.Write(context.Background(), c, protocol.Envelope{Type: "future.thing"}))
|
|
send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: protocol.Service{Kind: protocol.KindContainer, ContainerName: "y"}})
|
|
|
|
require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.upserts) })() == 1 }, time.Second, 5*time.Millisecond)
|
|
}
|
|
|
|
func TestHub_LargeResultArrives(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
store.finishOK = true
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
|
|
out := strings.Repeat("\x1b", protocol.MaxJobOutput)
|
|
sendJob(t, c, protocol.TypeJobResult, "42", protocol.JobResult{Status: protocol.ResultFailed, Output: out})
|
|
assert.Equal(t, "42", readEnv(t, c).ID)
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
require.Len(t, store.finishes, 1)
|
|
assert.Equal(t, out, store.finishes[0].Output)
|
|
}
|
|
|
|
func TestHub_ConcurrentWritesAreSerialised(t *testing.T) {
|
|
const n = 50
|
|
store := newFakeStore(node2)
|
|
store.finishOK = true
|
|
h, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
send(t, c, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
|
|
r := waitRoute(t, h)
|
|
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < n; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
assert.NoError(t, h.StartJob(r, int64(1000+i), protocol.JobStart{Ref: "r:1"}))
|
|
}()
|
|
}
|
|
go func() {
|
|
for i := 0; i < n; i++ {
|
|
sendJob(t, c, protocol.TypeJobResult, strconv.Itoa(i), protocol.JobResult{Status: protocol.ResultFailed})
|
|
}
|
|
}()
|
|
|
|
starts, acks := 0, 0
|
|
for starts+acks < 2*n {
|
|
switch env := readEnv(t, c); env.Type {
|
|
case protocol.TypeJobStart:
|
|
starts++
|
|
case protocol.TypeJobAck:
|
|
acks++
|
|
default:
|
|
t.Fatalf("unexpected %q", env.Type)
|
|
}
|
|
}
|
|
wg.Wait()
|
|
assert.Equal(t, n, starts)
|
|
assert.Equal(t, n, acks)
|
|
}
|