orbit/internal/transport/httpserver/event_hub_test.go

package httpserver

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/orbit/internal/jsonwire"
	core "git.yuntuo.hk/k2safe/orbit/model"
	"git.yuntuo.hk/k2safe/orbit/platform/auth"
	"git.yuntuo.hk/k2safe/orbit/platform/events"
	"github.com/coder/websocket"
)

type slowEventStore struct {
	events.Store
	once     sync.Once
	started  chan struct{}
	finished chan struct{}
}

func (s *slowEventStore) AuthorizeEvents(context.Context, auth.Token) error { return nil }
func (s *slowEventStore) BindConsumer(_ context.Context, _ auth.Token, id core.ID) (events.ConsumerBinding, error) {
	return events.ConsumerBinding{ApplicationID: core.ID(strings.Repeat("1", 32)), SubscriptionID: id, QueueID: core.ID(strings.Repeat("2", 32)), Enabled: true, ServerTime: time.Now().UTC()}, nil
}
func (s *slowEventStore) ClaimDelivery(ctx context.Context, _ auth.Token, _ core.ID, _ events.ClaimInput) (*events.Grant, error) {
	s.once.Do(func() { close(s.started) })
	<-ctx.Done()
	select {
	case s.finished <- struct{}{}:
	default:
	}
	return nil, ctx.Err()
}
func TestEventHubSlowClaimKeepsControlLiveAndBoundsConnections(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	store := &slowEventStore{started: make(chan struct{}), finished: make(chan struct{}, 4)}
	hub := NewEventHub(ctx, events.New(store))
	if cap(hub.slots) != 128 {
		t.Fatal("unexpected production connection bound")
	}
	hub.slots = make(chan struct{}, 2)
	hub.poll = 10 * time.Millisecond
	server := New("", nil)
	RegisterEventConsumers(server, hub)
	tls := httptest.NewTLSServer(server.Handler)
	defer tls.Close()
	defer hub.Shutdown(context.Background())
	id := core.ID(strings.Repeat("3", 32))
	credential := core.ID(strings.Repeat("4", 32))
	token, _, err := auth.Issue(id, credential)
	if err != nil {
		t.Fatal(err)
	}
	headers := http.Header{}
	headers.Set("Authorization", "Bearer "+token)
	dial := func() (*websocket.Conn, *http.Response, error) {
		return websocket.Dial(ctx, strings.Replace(tls.URL, "https://", "wss://", 1)+"/v1/events/connect", &websocket.DialOptions{HTTPClient: tls.Client(), HTTPHeader: headers, Subprotocols: []string{EventSubprotocol}})
	}
	connect := func() (*websocket.Conn, core.ID) {
		t.Helper()
		conn, _, err := dial()
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { conn.CloseNow() })
		hello, _ := jsonwire.Marshal(eventHello{SubscriptionID: id, Capacity: 1})
		raw, _ := jsonwire.Marshal(eventFrame{Version: "1", Type: "events.hello", MessageID: id, Body: hello})
		if err = conn.Write(ctx, websocket.MessageText, raw); err != nil {
			t.Fatal(err)
		}
		_, reply, err := conn.Read(ctx)
		if err != nil {
			t.Fatal(err)
		}
		var frame eventFrame
		if json.Unmarshal(reply, &frame) != nil || frame.Type != "events.welcome" {
			t.Fatal("missing welcome")
		}
		return conn, frame.SessionID
	}
	first, session := connect()
	second, _ := connect()
	select {
	case <-store.started:
	case <-ctx.Done():
		t.Fatal("claim was not dispatched")
	}
	blocked, resp, err := dial()
	if blocked != nil {
		blocked.CloseNow()
	}
	if err == nil || resp == nil || resp.StatusCode != 429 || resp.Header.Get("Retry-After") != "1" {
		t.Fatal("event connection bound not enforced")
	}
	ping, _ := jsonwire.Marshal(eventFrame{Version: "1", Type: "events.ping", MessageID: credential, SessionID: session, Body: json.RawMessage(`{}`)})
	if err = first.Write(ctx, websocket.MessageText, ping); err != nil {
		t.Fatal(err)
	}
	bounded, stop := context.WithTimeout(ctx, 500*time.Millisecond)
	_, raw, err := first.Read(bounded)
	stop()
	if err != nil {
		t.Fatal("slow claim blocked control reply")
	}
	var reply eventFrame
	if json.Unmarshal(raw, &reply) != nil || reply.Type != "ack" || reply.ReplyTo != credential {
		t.Fatal("missing ping acknowledgement")
	}
	second.CloseNow()
	deadline := time.Now().Add(time.Second)
	for len(hub.slots) == 2 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if len(hub.slots) != 1 {
		t.Fatal("disconnected consumer leaked connection slot")
	}
	third, _ := connect()
	defer third.CloseNow()
	bounded, stop = context.WithTimeout(ctx, time.Second)
	defer stop()
	if err = hub.Shutdown(bounded); err != nil {
		t.Fatal("shutdown leaked claim or connection", err)
	}
	select {
	case <-store.finished:
	default:
		t.Fatal("shutdown did not cancel in-flight claim")
	}
	if len(hub.slots) != 0 {
		t.Fatal("shutdown leaked connection slots")
	}
}