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()) } }