orbit/sdk/go/event_consumer_test.go

package orbit

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

	"github.com/coder/websocket"
)

func runEventConsumer(t *testing.T, client *Client, handle func(context.Context, ResultEvent) error) (*EventConsumer, context.CancelFunc, <-chan error) {
	t.Helper()
	consumer, err := NewEventConsumer(client, EventConsumerOptions{SubscriptionID: eventTestSub, Capacity: 1, Handle: handle, ReconnectMinDelay: 10 * time.Millisecond, ReconnectMaxDelay: 50 * time.Millisecond, ShutdownTimeout: time.Second})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- consumer.Run(ctx) }()
	t.Cleanup(cancel)
	return consumer, cancel, done
}
func awaitConsumerEvent(t *testing.T, c *EventConsumer, typ string) EventConsumerEvent {
	t.Helper()
	timer := time.NewTimer(3 * time.Second)
	defer timer.Stop()
	for {
		select {
		case event, ok := <-c.Events():
			if !ok {
				t.Fatal("consumer stopped before " + typ)
			}
			if event.Type == typ {
				return event
			}
		case <-timer.C:
			t.Fatal("consumer did not emit " + typ)
		}
	}
}
func awaitConsumerStop(t *testing.T, done <-chan error) error {
	t.Helper()
	select {
	case err := <-done:
		return err
	case <-time.After(3 * time.Second):
		t.Fatal("consumer did not stop")
		return nil
	}
}
func TestEventConsumerReconcilesLostAckWithoutRerunningHandler(t *testing.T) {
	g := eventTestGrant()
	var calls atomic.Int64
	var proofs [2]EventProof
	client := eventMock(t, func(ctx context.Context, c *websocket.Conn, _ EventRegistration, n int) {
		if n > 2 {
			t.Error("unexpected third connection")
			return
		}
		if n == 1 {
			_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(g)})
		}
		f, err := eventTestRead(ctx, c)
		if err != nil {
			return
		}
		if f.Type != "event.ack" || json.Unmarshal(f.Body, &proofs[n-1]) != nil {
			t.Error("missing ACK proof")
			return
		}
		if n == 1 {
			return
		} // ACK已在服务端提交，回应在网络中丢失。
		if proofs[0].DeliveryID != proofs[1].DeliveryID || proofs[0].Epoch != proofs[1].Epoch || proofs[0].LeaseSecret != proofs[1].LeaseSecret {
			t.Error("reconnect changed acknowledgement identity")
		}
		out := eventAckFor(g)
		out.Deduplicated = true
		_ = eventTestWrite(ctx, c, eventFrame{Type: "ack", ReplyTo: f.MessageID, Body: eventBody(out)})
		_, _ = eventTestRead(ctx, c)
	})
	consumer, cancel, done := runEventConsumer(t, client, func(context.Context, ResultEvent) error { calls.Add(1); return nil })
	awaitConsumerEvent(t, consumer, "acknowledged")
	cancel()
	if err := awaitConsumerStop(t, done); err != nil {
		t.Fatal(err)
	}
	if calls.Load() != 1 || consumer.Stats().Acknowledged != 1 || consumer.Stats().Reconnects != 1 {
		t.Fatal("lost receipt re-executed committed handler", consumer.Stats())
	}
	if !errors.Is(consumer.Run(context.Background()), ErrEventConsumerAlreadyRun) {
		t.Fatal("consumer ran twice")
	}
}
func TestEventConsumerSupersededProofRequiresDurableDeduplication(t *testing.T) {
	g := eventTestGrant()
	var calls, effects atomic.Int64
	client := eventMock(t, func(ctx context.Context, c *websocket.Conn, _ EventRegistration, _ int) {
		_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(g)})
		f, err := eventTestRead(ctx, c)
		if err != nil {
			return
		}
		_ = eventTestWrite(ctx, c, eventFrame{Type: "error", ReplyTo: f.MessageID, Error: &EventError{Code: "DELIVERY_LEASE_LOST", Message: "proof superseded"}})
		next := g
		next.Delivery.LeaseEpoch = 2
		next.Delivery.Attempts = 2
		next.Delivery.Revision = 4
		_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(next)})
		f, err = eventTestRead(ctx, c)
		if err != nil {
			return
		}
		_ = eventTestWrite(ctx, c, eventFrame{Type: "ack", ReplyTo: f.MessageID, Body: eventBody(eventAckFor(next))})
		_, _ = eventTestRead(ctx, c)
	})
	consumer, cancel, done := runEventConsumer(t, client, func(context.Context, ResultEvent) error { calls.Add(1); effects.CompareAndSwap(0, 1); return nil })
	awaitConsumerEvent(t, consumer, "acknowledged")
	cancel()
	if err := awaitConsumerStop(t, done); err != nil {
		t.Fatal(err)
	}
	if calls.Load() != 2 || effects.Load() != 1 {
		t.Fatal("replacement proof bypassed handler deduplication")
	}
}
func TestEventConsumerHandlerFailureDoesNotAcknowledge(t *testing.T) {
	for _, panicHandler := range []bool{false, true} {
		name := "error"
		if panicHandler {
			name = "panic"
		}
		t.Run(name, func(t *testing.T) {
			g := eventTestGrant()
			retry := make(chan struct{})
			var calls atomic.Int64
			client := eventMock(t, func(ctx context.Context, c *websocket.Conn, _ EventRegistration, _ int) {
				_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(g)})
				select {
				case <-retry:
				case <-ctx.Done():
					return
				}
				next := g
				next.Delivery.LeaseEpoch = 2
				next.Delivery.Attempts = 2
				_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(next)})
				f, err := eventTestRead(ctx, c)
				if err != nil {
					return
				}
				var proof EventProof
				if f.Type != "event.ack" || json.Unmarshal(f.Body, &proof) != nil || proof.Epoch != 2 {
					t.Error("failed handler acknowledged original delivery")
				}
				_ = eventTestWrite(ctx, c, eventFrame{Type: "ack", ReplyTo: f.MessageID, Body: eventBody(eventAckFor(next))})
				_, _ = eventTestRead(ctx, c)
			})
			consumer, cancel, done := runEventConsumer(t, client, func(context.Context, ResultEvent) error {
				if calls.Add(1) == 1 {
					if panicHandler {
						panic("private handler data")
					}
					return errors.New("private handler data")
				}
				return nil
			})
			awaitConsumerEvent(t, consumer, "handler_failed")
			if calls.Load() != 1 {
				t.Fatal("handler retried locally")
			}
			close(retry)
			awaitConsumerEvent(t, consumer, "acknowledged")
			cancel()
			if err := awaitConsumerStop(t, done); err != nil {
				t.Fatal(err)
			}
			if consumer.Stats().HandlerFailures != 1 || consumer.Stats().Acknowledged != 1 {
				t.Fatal(consumer.Stats())
			}
		})
	}
}
func TestEventConsumerShutdownCancelsAndConfirmsCommittedHandler(t *testing.T) {
	g := eventTestGrant()
	started := make(chan struct{})
	client := eventMock(t, func(ctx context.Context, c *websocket.Conn, _ EventRegistration, _ int) {
		_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(g)})
		f, err := eventTestRead(ctx, c)
		if err != nil {
			return
		}
		_ = eventTestWrite(ctx, c, eventFrame{Type: "ack", ReplyTo: f.MessageID, Body: eventBody(eventAckFor(g))})
		_, _ = eventTestRead(ctx, c)
	})
	consumer, cancel, done := runEventConsumer(t, client, func(ctx context.Context, _ ResultEvent) error { close(started); <-ctx.Done(); return nil })
	<-started
	cancel()
	if err := awaitConsumerStop(t, done); err != nil {
		t.Fatal(err)
	}
	if consumer.Stats().Acknowledged != 1 || consumer.Stats().Running != 0 {
		t.Fatal("committed result not drained", consumer.Stats())
	}
}
func TestEventConsumerShutdownBoundsIgnoringHandler(t *testing.T) {
	started, release := make(chan struct{}), make(chan struct{})
	defer close(release)
	client := eventMock(t, func(ctx context.Context, c *websocket.Conn, _ EventRegistration, _ int) {
		_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(eventTestGrant())})
		_, _ = eventTestRead(ctx, c)
	})
	consumer, err := NewEventConsumer(client, EventConsumerOptions{SubscriptionID: eventTestSub, Capacity: 1, ShutdownTimeout: 30 * time.Millisecond, Handle: func(context.Context, ResultEvent) error { close(started); <-release; return nil }})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- consumer.Run(ctx) }()
	<-started
	cancel()
	err = awaitConsumerStop(t, done)
	var run *EventConsumerRunError
	if !errors.Is(err, ErrEventShutdownTimeout) || !errors.As(err, &run) || run.Unconfirmed != 1 || run.Running != 1 {
		t.Fatal("uncooperative handler not reported", err)
	}
}
func TestEventConsumerPermanentFailureDoesNotReconnect(t *testing.T) {
	var calls atomic.Int64
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1); w.WriteHeader(http.StatusUnauthorized) }))
	defer server.Close()
	client, _ := NewClient(server.URL, Options{Token: "token"})
	_, _, done := runEventConsumer(t, client, func(context.Context, ResultEvent) error { t.Error("unauthorized consumer invoked handler"); return nil })
	err := awaitConsumerStop(t, done)
	var status *HTTPError
	if !errors.As(err, &status) || status.StatusCode != 401 || calls.Load() != 1 {
		t.Fatal("permanent handshake error retried", err)
	}
}

func TestEventConsumerHandlerGoexitReleasesSlotWithoutAck(t *testing.T) {
	client := eventMock(t, func(ctx context.Context, c *websocket.Conn, _ EventRegistration, _ int) {
		_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(eventTestGrant())})
		if f, err := eventTestRead(ctx, c); err == nil && f.Type == "event.ack" {
			t.Error("Goexit handler was acknowledged")
		}
	})
	consumer, cancel, done := runEventConsumer(t, client, func(context.Context, ResultEvent) error { runtime.Goexit(); return nil })
	awaitConsumerEvent(t, consumer, "handler_failed")
	cancel()
	if err := awaitConsumerStop(t, done); err != nil || consumer.Stats().Running != 0 {
		t.Fatal("Goexit leaked handler capacity", err)
	}
}

func TestEventConsumerConcurrentCapacity(t *testing.T) {
	start := make(chan string, 3)
	release := make(chan struct{})
	var active, peak atomic.Int64
	grants := []EventGrant{eventTestGrant(), eventTestGrant(), eventTestGrant()}
	for i := range grants {
		grants[i].Event.ID = workerTestID(t)
		grants[i].Delivery.ID = workerTestID(t)
		grants[i].Delivery.EventID = grants[i].Event.ID
	}
	client := eventMock(t, func(ctx context.Context, c *websocket.Conn, in EventRegistration, _ int) {
		if in.Capacity != 2 {
			t.Error("capacity changed")
		}
		for _, g := range grants[:2] {
			_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(g)})
		}
		for i := 0; i < 3; i++ {
			f, err := eventTestRead(ctx, c)
			if err != nil {
				return
			}
			var p EventProof
			if f.Type != "event.ack" || json.Unmarshal(f.Body, &p) != nil {
				t.Error("missing ACK")
				return
			}
			var grant EventGrant
			for _, g := range grants {
				if g.Delivery.ID == p.DeliveryID {
					grant = g
				}
			}
			_ = eventTestWrite(ctx, c, eventFrame{Type: "ack", ReplyTo: f.MessageID, Body: eventBody(eventAckFor(grant))})
			if i == 0 {
				_ = eventTestWrite(ctx, c, eventFrame{Type: "event.deliver", Body: eventBody(grants[2])})
			}
		}
		_, _ = eventTestRead(ctx, c)
	})
	consumer, err := NewEventConsumer(client, EventConsumerOptions{SubscriptionID: eventTestSub, Capacity: 2, Handle: func(ctx context.Context, event ResultEvent) error {
		n := active.Add(1)
		defer active.Add(-1)
		for {
			old := peak.Load()
			if n <= old || peak.CompareAndSwap(old, n) {
				break
			}
		}
		start <- event.ID
		select {
		case <-release:
			return nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- consumer.Run(ctx) }()
	for range 2 {
		select {
		case <-start:
		case <-ctx.Done():
			t.Fatal("capacity not utilized")
		}
	}
	close(release)
	for range 3 {
		awaitConsumerEvent(t, consumer, "acknowledged")
	}
	cancel()
	if err := awaitConsumerStop(t, done); err != nil || peak.Load() != 2 || consumer.Stats().Started != 3 {
		t.Fatal("capacity or drain invariant", err, consumer.Stats())
	}
}