From a87f4ca66dd5e792d56d1cf085ed3559ce513d7a Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Wed, 11 Nov 2026 10:31:49 +1100 Subject: [PATCH 1/5] Failure events: agent detects lost invocations and the Go SDK hooks them The agent now reports failures it sees on its own (queue full, expired while queued, no report before the timeout) on a new DispatchEvents stream, and re-queues a failed invocation on RetryInvocation. The queue behind the stream is unbounded so no failure is dropped, and it only exists for clients that set DispatchRequest.failure_events, so older SDKs are unaffected. The Go SDK adds WithPreDispatchHook and WithPostDispatchHook. The post hook runs for every invocation, plus every agent failure event, with a reason (HANDLER_ERROR, TIMEOUT, NO_RESOURCES) and a source (APPLICATION or RUNTIME), and can ask for a retry. A timeout while a Cortex API call is in flight is classed RUNTIME. invokeHandler no longer lets a timed-out handler write to the report while it is being sent. Co-Authored-By: Claude Opus 5.5 --- agent/proto/cortex-axon-agent.proto | 60 ++ agent/server/failure_events_e2e_test.go | 208 ++++++ agent/server/handler/failure_events_test.go | 194 ++++++ agent/server/handler/failure_queue.go | 78 +++ agent/server/handler/handler_invoke.go | 24 + agent/server/handler/handler_manager_test.go | 3 +- agent/server/handler/manager.go | 124 ++++ .../handler/scheduled_handler_entry_test.go | 20 + agent/server/http/http_server_webhook_test.go | 11 + agent/server/server.go | 97 ++- agent/server/test_axon_client.go | 21 +- .../cortexapps/axon/cortex-axon-agent.pb.go | 600 ++++++++++++++---- .../axon/cortex-axon-agent_grpc.pb.go | 91 ++- sdks/go/axon_agent.go | 108 +++- sdks/go/axon_agent_test.go | 29 +- sdks/go/handler_context.go | 37 +- sdks/go/hooks.go | 203 ++++++ sdks/go/hooks_test.go | 284 +++++++++ sdks/go/mock_axon/mocks.go | 40 ++ sdks/go/options.go | 2 + 20 files changed, 2078 insertions(+), 156 deletions(-) create mode 100644 agent/server/failure_events_e2e_test.go create mode 100644 agent/server/handler/failure_events_test.go create mode 100644 agent/server/handler/failure_queue.go create mode 100644 sdks/go/hooks.go create mode 100644 sdks/go/hooks_test.go diff --git a/agent/proto/cortex-axon-agent.proto b/agent/proto/cortex-axon-agent.proto index 760e5a7f..09dd16b8 100644 --- a/agent/proto/cortex-axon-agent.proto +++ b/agent/proto/cortex-axon-agent.proto @@ -12,6 +12,15 @@ service AxonAgent { rpc GetHandlerHistory(GetHandlerHistoryRequest) returns (GetHandlerHistoryResponse); rpc Dispatch (stream DispatchRequest) returns (stream DispatchMessage); rpc ReportInvocation(ReportInvocationRequest) returns (ReportInvocationResponse); + + // DispatchEvents streams failures the agent detected on its own: invocations + // that never reached the client, or that the client never reported. They use + // their own stream so a client that has stopped reading Dispatch because + // every handler is busy still hears about them. + rpc DispatchEvents(DispatchRequest) returns (stream DispatchEvent); + + // RetryInvocation queues a failed invocation again as a new attempt. + rpc RetryInvocation(RetryInvocationRequest) returns (RetryInvocationResponse); } @@ -86,6 +95,9 @@ message ListHandlersResponse { message DispatchRequest { string dispatch_id = 1; string client_version = 2; + // The client handles failure events. The agent only keeps them for clients + // that ask, so an older client never has events pile up for it. + bool failure_events = 3; } message DispatchMessage { @@ -111,6 +123,54 @@ message DispatchHandlerInvoke{ int32 timeout_ms = 10; HandlerInvokeType reason = 11; map args = 20; + // 1 for the first attempt; 0 means the agent didn't say, so treat it as 1. + int32 attempt = 30; +} + +// +// Failure events +// + +enum FailureReason { + FAILURE_REASON_UNSPECIFIED = 0; + // The handler returned an error or raised. + HANDLER_ERROR = 1; + // The invocation did not finish within its timeout. + TIMEOUT = 2; + // The invocation could not be dispatched: the queue was full, or it expired + // while waiting for a free handler. + NO_RESOURCES = 3; +} + +enum FailureSource { + FAILURE_SOURCE_UNSPECIFIED = 0; + // The failure came from handler code. + APPLICATION = 1; + // The failure came from Axon: the agent, the SDK, or a Cortex API call. + RUNTIME = 2; +} + +message InvocationFailure { + FailureReason reason = 1; + FailureSource source = 2; + string message = 3; + // Whoever submitted the invocation was told it was not accepted (e.g. a + // webhook sender got a 503), so it may be sent again. + bool caller_rejected = 4; +} + +message DispatchEvent { + DispatchHandlerInvoke invoke = 1; + InvocationFailure failure = 2; +} + +message RetryInvocationRequest { + DispatchHandlerInvoke invoke = 1; +} + +message RetryInvocationResponse { + Error error = 1; + string invocation_id = 2; } // diff --git a/agent/server/failure_events_e2e_test.go b/agent/server/failure_events_e2e_test.go new file mode 100644 index 00000000..a9b8d7f9 --- /dev/null +++ b/agent/server/failure_events_e2e_test.go @@ -0,0 +1,208 @@ +package server + +import ( + "context" + "errors" + "net/url" + "path/filepath" + "sync" + "testing" + "time" + + pb "github.com/cortexapps/axon/.generated/proto/github.com/cortexapps/axon" + "github.com/cortexapps/axon/config" + "github.com/cortexapps/axon/server/cron" + "github.com/cortexapps/axon/server/handler" + "github.com/stretchr/testify/require" + "go.uber.org/zap" + "google.golang.org/protobuf/proto" +) + +const failureTestWebhook = "failure-test-hook" + +type failureEventsFixture struct { + t *testing.T + ctx context.Context + manager handler.Manager + client *testAxonClient + + mu sync.Mutex + invokes []*pb.DispatchHandlerInvoke + received chan *pb.DispatchHandlerInvoke +} + +// startFailureEventsFixture runs an agent and a client with one webhook +// handler. respond decides what the client does with each invocation. +func startFailureEventsFixture(t *testing.T, timeout time.Duration, respond func(*pb.DispatchHandlerInvoke) (clientResult, error)) *failureEventsFixture { + oldGrace := reportGrace + reportGrace = 100 * time.Millisecond + t.Cleanup(func() { reportGrace = oldGrace }) + + port := getRandomPort() + cfg := config.AgentConfig{ + GrpcPort: port, + CortexApiBaseUrl: "http://localhost", + CortexApiToken: "test-token", + DequeueWaitTime: 100 * time.Millisecond, + HandlerHistoryPath: filepath.Join(t.TempDir(), "history"), + } + + logger, _ := zap.NewDevelopment() + manager := handler.NewHandlerManager(logger, cron.New(), nil) + agent := NewAxonAgent(Params{Logger: logger, Config: cfg, Manager: manager}) + require.NoError(t, agent.Start(context.Background())) + t.Cleanup(agent.Close) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + t.Cleanup(cancel) + + f := &failureEventsFixture{ + t: t, + ctx: ctx, + manager: manager, + received: make(chan *pb.DispatchHandlerInvoke, 10), + } + + f.client = NewTestAxonClient(t, int32(port), func(msg *pb.DispatchMessage) (clientResult, error) { + invoke := msg.GetInvoke() + if invoke == nil { + return clientResult{}, nil + } + f.received <- invoke + return respond(invoke) + }) + f.client.failureEvents = true + t.Cleanup(f.client.Close) + + _, err := f.client.client.RegisterHandler(ctx, &pb.RegisterHandlerRequest{ + DispatchId: f.client.dispatchId, + HandlerName: "hook", + TimeoutMs: int32(timeout.Milliseconds()), + Options: []*pb.HandlerOption{{ + Option: &pb.HandlerOption_Invoke{ + Invoke: &pb.HandlerInvokeOption{Type: pb.HandlerInvokeType_WEBHOOK, Value: failureTestWebhook}, + }, + }}, + }) + require.NoError(t, err) + + go f.client.Run(ctx) + require.Eventually(t, func() bool { return manager.GetByTag(failureTestWebhook) != nil }, + 5*time.Second, 10*time.Millisecond, "handler never started") + return f +} + +func (f *failureEventsFixture) postWebhook(body string) { + u, _ := url.Parse("http://localhost/webhook/" + failureTestWebhook) + entry := f.manager.GetByTag(failureTestWebhook) + require.NoError(f.t, f.manager.Trigger(handler.NewWebhookHandlerInvoke(entry, u, body, "application/json"))) +} + +func (f *failureEventsFixture) nextInvoke() *pb.DispatchHandlerInvoke { + select { + case invoke := <-f.received: + return invoke + case <-time.After(5 * time.Second): + require.FailNow(f.t, "client never received an invocation") + return nil + } +} + +func (f *failureEventsFixture) events() <-chan *pb.DispatchEvent { + stream, err := f.client.DispatchEvents(f.ctx) + require.NoError(f.t, err) + out := make(chan *pb.DispatchEvent, 10) + go func() { + for { + ev, err := stream.Recv() + if err != nil { + return + } + out <- ev + } + }() + return out +} + +func requireEvent(t *testing.T, events <-chan *pb.DispatchEvent) *pb.DispatchEvent { + select { + case ev := <-events: + return ev + case <-time.After(5 * time.Second): + require.FailNow(t, "no failure event arrived") + return nil + } +} + +func requireNoEvent(t *testing.T, events <-chan *pb.DispatchEvent, wait time.Duration) { + select { + case ev := <-events: + require.FailNow(t, "unexpected failure event", "%v", ev) + case <-time.After(wait): + } +} + +func TestFailureEvents_UnreportedInvocationTimesOutAndRetries(t *testing.T) { + calls := 0 + f := startFailureEventsFixture(t, 200*time.Millisecond, func(*pb.DispatchHandlerInvoke) (clientResult, error) { + calls++ + if calls == 1 { + return clientResult{noReport: true}, nil + } + return clientResult{result: "ok"}, nil + }) + + // Connect after the failure so delivery doesn't depend on a live stream. + f.postWebhook(`{"id": 1}`) + first := f.nextInvoke() + require.Equal(t, int32(1), first.Attempt) + time.Sleep(500 * time.Millisecond) + + events := f.events() + ev := requireEvent(t, events) + require.Equal(t, first.InvocationId, ev.Invoke.InvocationId) + require.Equal(t, `{"id": 1}`, ev.Invoke.Args["body"]) + require.Equal(t, pb.FailureReason_TIMEOUT, ev.Failure.Reason) + require.Equal(t, pb.FailureSource_RUNTIME, ev.Failure.Source) + + resp, err := f.client.RetryInvocation(f.ctx, ev.Invoke) + require.NoError(t, err) + require.Nil(t, resp.Error) + + retried := f.nextInvoke() + require.Equal(t, resp.InvocationId, retried.InvocationId) + require.Equal(t, int32(2), retried.Attempt) + require.Equal(t, `{"id": 1}`, retried.Args["body"]) + + // The retry succeeded, so nothing else fails. + requireNoEvent(t, events, time.Second) +} + +func TestFailureEvents_ClientReportedFailureIsNotRepeated(t *testing.T) { + f := startFailureEventsFixture(t, 200*time.Millisecond, func(*pb.DispatchHandlerInvoke) (clientResult, error) { + // The client saw the failure and fires its own hook. + return clientResult{}, errors.New("handler blew up") + }) + events := f.events() + + f.postWebhook(`{}`) + f.nextInvoke() + + // Past the timeout and grace, the agent must stay quiet. + requireNoEvent(t, events, time.Second) +} + +func TestFailureEvents_RetryRejectedAtLimit(t *testing.T) { + f := startFailureEventsFixture(t, time.Second, func(*pb.DispatchHandlerInvoke) (clientResult, error) { + return clientResult{result: "ok"}, nil + }) + + f.postWebhook(`{}`) + invoke := proto.Clone(f.nextInvoke()).(*pb.DispatchHandlerInvoke) + invoke.Attempt = 5 + + resp, err := f.client.RetryInvocation(f.ctx, invoke) + require.NoError(t, err) + require.NotNil(t, resp.Error) + require.Equal(t, "rejected", resp.Error.Code) +} diff --git a/agent/server/handler/failure_events_test.go b/agent/server/handler/failure_events_test.go new file mode 100644 index 00000000..c04bdf1b --- /dev/null +++ b/agent/server/handler/failure_events_test.go @@ -0,0 +1,194 @@ +package handler + +import ( + "context" + "fmt" + "testing" + "time" + + pb "github.com/cortexapps/axon/.generated/proto/github.com/cortexapps/axon" + "github.com/cortexapps/axon/server/cron" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +func webhookOnlyOption() *pb.HandlerOption { + return &pb.HandlerOption{ + Option: &pb.HandlerOption_Invoke{ + Invoke: &pb.HandlerInvokeOption{ + Type: pb.HandlerInvokeType_WEBHOOK, + }, + }, + } +} + +func startWebhookHandler(t *testing.T, timeout time.Duration) (Manager, HandlerEntry) { + logger, _ := zap.NewDevelopment() + mgr := NewHandlerManager(logger, cron.New(), nil) + _, err := mgr.RegisterHandler("1", "handler1", timeout, webhookOnlyOption()) + require.NoError(t, err) + require.NoError(t, mgr.Start("1")) + return mgr, mgr.ListHandlers()[0] +} + +func nextFailure(t *testing.T, mgr Manager) *pb.DispatchEvent { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + ev, err := mgr.NextFailure(ctx, "1") + require.NoError(t, err, "expected a failure event") + return ev +} + +func requireNoFailure(t *testing.T, mgr Manager) { + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + ev, err := mgr.NextFailure(ctx, "1") + require.Error(t, err, "unexpected failure event: %v", ev) +} + +func TestQueueFullEmitsNoResources(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + mgr.SubscribeFailures("1") + + for i := 0; i < dispatchQueueDepth; i++ { + require.NoError(t, mgr.Trigger(NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil))) + } + + rejected := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"body": "payload"}) + require.ErrorIs(t, mgr.Trigger(rejected), ErrDispatchQueueFull) + + ev := nextFailure(t, mgr) + require.Equal(t, rejected.Id, ev.Invoke.InvocationId) + require.Equal(t, "payload", ev.Invoke.Args["body"]) + require.Equal(t, pb.FailureReason_NO_RESOURCES, ev.Failure.Reason) + require.Equal(t, pb.FailureSource_RUNTIME, ev.Failure.Source) + require.True(t, ev.Failure.CallerRejected) +} + +func TestExpiredInQueueEmitsNoResources(t *testing.T) { + mgr, entry := startWebhookHandler(t, 50*time.Millisecond) + mgr.SubscribeFailures("1") + + invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil) + require.NoError(t, mgr.Trigger(invoke)) + + // Nothing dequeues, so it expires in the queue. + ev := nextFailure(t, mgr) + require.Equal(t, invoke.Id, ev.Invoke.InvocationId) + require.Equal(t, pb.FailureReason_NO_RESOURCES, ev.Failure.Reason) + require.False(t, ev.Failure.CallerRejected) + + // The dispatch loop must not send it once it has expired. + h, err := mgr.Dequeue(context.Background(), "1", time.Second) + require.NoError(t, err) + require.NotNil(t, h) + require.False(t, h.ClaimDispatch()) +} + +func TestDispatchedInvocationDoesNotEmitOnExpiry(t *testing.T) { + mgr, entry := startWebhookHandler(t, 50*time.Millisecond) + mgr.SubscribeFailures("1") + + require.NoError(t, mgr.Trigger(NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil))) + h, err := mgr.Dequeue(context.Background(), "1", time.Second) + require.NoError(t, err) + require.True(t, h.ClaimDispatch()) + + // The agent's report timer owns timeouts for dispatched work. + <-h.Done() + requireNoFailure(t, mgr) +} + +func TestFailuresWithoutSubscriberAreNotQueued(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + + mgr.EmitFailure(NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke(), + &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) + + // Subscribing later must not surface the earlier event. + requireNoFailure(t, mgr) +} + +func TestFailuresAreNeverDropped(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + mgr.SubscribeFailures("1") + + const count = 10000 + for i := 0; i < count; i++ { + invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"n": fmt.Sprint(i)}) + mgr.EmitFailure(invoke.ToDispatchInvoke(), &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) + } + + for i := 0; i < count; i++ { + ev := nextFailure(t, mgr) + require.Equal(t, fmt.Sprint(i), ev.Invoke.Args["n"], "events must arrive in order") + } + requireNoFailure(t, mgr) +} + +func TestRequeuedFailureIsDeliveredNext(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + mgr.SubscribeFailures("1") + + for _, n := range []string{"a", "b"} { + invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"n": n}) + mgr.EmitFailure(invoke.ToDispatchInvoke(), &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) + } + + first := nextFailure(t, mgr) + require.Equal(t, "a", first.Invoke.Args["n"]) + mgr.RequeueFailure("1", first) + + require.Equal(t, "a", nextFailure(t, mgr).Invoke.Args["n"]) + require.Equal(t, "b", nextFailure(t, mgr).Invoke.Args["n"]) +} + +func TestRetriggerQueuesNextAttempt(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + + original := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"body": "payload"}) + id, err := mgr.Retrigger(original.ToDispatchInvoke()) + require.NoError(t, err) + require.NotEqual(t, original.Id, id) + + h, err := mgr.Dequeue(context.Background(), "1", time.Second) + require.NoError(t, err) + retried := h.ToDispatchInvoke() + require.Equal(t, id, retried.InvocationId) + require.Equal(t, int32(2), retried.Attempt) + require.Equal(t, "payload", retried.Args["body"]) + require.Equal(t, pb.HandlerInvokeType_WEBHOOK, retried.Reason) +} + +func TestRetriggerFindsReregisteredHandler(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + + stale := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke() + stale.HandlerId = "gone" + + _, err := mgr.Retrigger(stale) + require.NoError(t, err) +} + +func TestRetriggerStopsAtLimit(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + + last := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke() + last.Attempt = maxRetryAttempts + + _, err := mgr.Retrigger(last) + require.ErrorIs(t, err, ErrRetryLimit) +} + +func TestRetriggerTreatsMissingAttemptAsFirst(t *testing.T) { + mgr, entry := startWebhookHandler(t, defaultTimeout) + + fromOlderAgent := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke() + fromOlderAgent.Attempt = 0 + + _, err := mgr.Retrigger(fromOlderAgent) + require.NoError(t, err) + h, err := mgr.Dequeue(context.Background(), "1", time.Second) + require.NoError(t, err) + require.Equal(t, int32(2), h.ToDispatchInvoke().Attempt) +} diff --git a/agent/server/handler/failure_queue.go b/agent/server/handler/failure_queue.go new file mode 100644 index 00000000..c8b717b9 --- /dev/null +++ b/agent/server/handler/failure_queue.go @@ -0,0 +1,78 @@ +package handler + +import ( + "context" + "sync" + + pb "github.com/cortexapps/axon/.generated/proto/github.com/cortexapps/axon" +) + +// failureQueue holds failure events for one dispatch id until its client +// reads them. It is unbounded on purpose: invocations are throttled, but a +// dropped failure is a message lost without a trace, which is the problem +// failure events exist to solve. +type failureQueue struct { + mu sync.Mutex + events []*pb.DispatchEvent + // notify has room for one wakeup, so push never blocks and a waiting + // pop never misses an event that arrives while it is between checks. + notify chan struct{} +} + +func newFailureQueue() *failureQueue { + return &failureQueue{notify: make(chan struct{}, 1)} +} + +func (q *failureQueue) push(ev *pb.DispatchEvent) { + q.mu.Lock() + q.events = append(q.events, ev) + q.mu.Unlock() + q.wake() +} + +// pushFront returns an event whose delivery failed, so it is the next one sent. +func (q *failureQueue) pushFront(ev *pb.DispatchEvent) { + q.mu.Lock() + q.events = append([]*pb.DispatchEvent{ev}, q.events...) + q.mu.Unlock() + q.wake() +} + +func (q *failureQueue) wake() { + select { + case q.notify <- struct{}{}: + default: + } +} + +// pop blocks until an event is available or ctx is done. +func (q *failureQueue) pop(ctx context.Context) (*pb.DispatchEvent, error) { + for { + q.mu.Lock() + if len(q.events) > 0 { + ev := q.events[0] + q.events[0] = nil + q.events = q.events[1:] + remaining := len(q.events) + q.mu.Unlock() + if remaining > 0 { + // Another reader may be waiting on the same queue. + q.wake() + } + return ev, nil + } + q.mu.Unlock() + + select { + case <-q.notify: + case <-ctx.Done(): + return nil, ctx.Err() + } + } +} + +func (q *failureQueue) len() int { + q.mu.Lock() + defer q.mu.Unlock() + return len(q.events) +} diff --git a/agent/server/handler/handler_invoke.go b/agent/server/handler/handler_invoke.go index 1bc3804a..b93f9c9c 100644 --- a/agent/server/handler/handler_invoke.go +++ b/agent/server/handler/handler_invoke.go @@ -19,8 +19,19 @@ type Invocable interface { Complete(result string, err error) error Done() <-chan struct{} GetResult() (string, error) + + // ClaimDispatch and ClaimExpired decide, exactly once, whether a queued + // invocation was sent to a client or expired before one picked it up. + ClaimDispatch() bool + ClaimExpired() bool } +const ( + invokeQueued int32 = iota + invokeDispatched + invokeExpired +) + type HandlerInvoke struct { Id string Entry HandlerEntry @@ -28,6 +39,7 @@ type HandlerInvoke struct { Args map[string]string Timeout time.Duration Timestamp time.Time + Attempt int32 result string err error @@ -35,6 +47,7 @@ type HandlerInvoke struct { started *atomic.Bool finished *atomic.Bool + state *atomic.Int32 } func NewHandlerInvoke( @@ -49,10 +62,12 @@ func NewHandlerInvoke( Args: args, Timeout: entry.Timeout(), Timestamp: time.Now(), + Attempt: 1, done: make(chan struct{}), started: &atomic.Bool{}, finished: &atomic.Bool{}, + state: &atomic.Int32{}, } } @@ -107,6 +122,14 @@ func (h HandlerInvoke) Done() <-chan struct{} { return h.done } +func (h HandlerInvoke) ClaimDispatch() bool { + return h.state.CompareAndSwap(invokeQueued, invokeDispatched) +} + +func (h HandlerInvoke) ClaimExpired() bool { + return h.state.CompareAndSwap(invokeQueued, invokeExpired) +} + func (h HandlerInvoke) ToDispatchInvoke() *pb.DispatchHandlerInvoke { return &pb.DispatchHandlerInvoke{ InvocationId: h.Id, @@ -116,5 +139,6 @@ func (h HandlerInvoke) ToDispatchInvoke() *pb.DispatchHandlerInvoke { Reason: h.Reason, Args: h.Args, TimeoutMs: int32(h.Entry.Timeout().Milliseconds()), + Attempt: h.Attempt, } } diff --git a/agent/server/handler/handler_manager_test.go b/agent/server/handler/handler_manager_test.go index 7bf7c2c7..b817b466 100644 --- a/agent/server/handler/handler_manager_test.go +++ b/agent/server/handler/handler_manager_test.go @@ -110,9 +110,10 @@ func TestTriggerAndDequeue(t *testing.T) { time.Sleep(time.Millisecond * 10) metrics, err := registry.Gather() require.NoError(t, err) - require.Equal(t, 3, len(metrics)) + require.Equal(t, 4, len(metrics)) expected := []string{ + "axon_handler_failure_events_pending", "axon_handler_invokes", "axon_handler_latency", "axon_handler_queue_depth", diff --git a/agent/server/handler/manager.go b/agent/server/handler/manager.go index 10358e10..93bb69a1 100644 --- a/agent/server/handler/manager.go +++ b/agent/server/handler/manager.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "os" + "sync" "time" pb "github.com/cortexapps/axon/.generated/proto/github.com/cortexapps/axon" @@ -22,6 +23,14 @@ const dispatchQueueDepth = 100 // response rather than accepting work that will only time out. var ErrDispatchQueueFull = errors.New("dispatch queue is full") +// maxRetryAttempts caps RetryInvocation. The failure hook sees the attempt +// number and normally decides when to stop; this only bounds a hook that +// always retries. +const maxRetryAttempts = 5 + +// ErrRetryLimit means an invocation has used all of its attempts. +var ErrRetryLimit = errors.New("retry limit reached") + type Manager interface { RegisterHandler(dispatchId string, name string, timeout time.Duration, options ...*pb.HandlerOption) (string, error) UnregisterHandler(id string) @@ -34,6 +43,18 @@ type Manager interface { Dequeue(ctx context.Context, id string, waitTime time.Duration) (Invocable, error) Close() error IsFinished() bool + + // SubscribeFailures starts keeping failure events for a dispatch id. + // Events for a dispatch id nobody subscribed are logged and discarded. + SubscribeFailures(dispatchId string) + EmitFailure(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) + // NextFailure blocks until a failure event is available or ctx is done. + NextFailure(ctx context.Context, dispatchId string) (*pb.DispatchEvent, error) + // RequeueFailure puts back an event that could not be delivered. + RequeueFailure(dispatchId string, ev *pb.DispatchEvent) + // Retrigger queues a failed invocation again as its next attempt and + // returns the new invocation id. + Retrigger(invoke *pb.DispatchHandlerInvoke) (string, error) } type handlerManager struct { @@ -47,6 +68,10 @@ type handlerManager struct { invokeCounter *prometheus.CounterVec queueDepthGauge *prometheus.GaugeVec handlerLatencyGauge *prometheus.HistogramVec + failuresPending prometheus.Gauge + + failureLock sync.Mutex + failureQueues map[string]*failureQueue } func NewHandlerManager(logger *zap.Logger, cron cron.Cron, metrics *prometheus.Registry) Manager { @@ -57,6 +82,7 @@ func NewHandlerManager(logger *zap.Logger, cron cron.Cron, metrics *prometheus.R dispatchQueues: make(map[string]chan Invocable), cron: cron, outstandingRequests: make(map[string]Invocable), + failureQueues: make(map[string]*failureQueue), invokeCounter: prometheus.NewCounterVec( prometheus.CounterOpts{ Name: "axon_handler_invokes", @@ -79,12 +105,19 @@ func NewHandlerManager(logger *zap.Logger, cron cron.Cron, metrics *prometheus.R }, []string{"handler", "result"}, ), + failuresPending: prometheus.NewGauge( + prometheus.GaugeOpts{ + Name: "axon_handler_failure_events_pending", + Help: "Number of failure events waiting for a client to read them", + }, + ), done: make(chan struct{}), } if metrics != nil { metrics.MustRegister(mgr.invokeCounter) metrics.MustRegister(mgr.queueDepthGauge) metrics.MustRegister(mgr.handlerLatencyGauge) + metrics.MustRegister(mgr.failuresPending) } return mgr } @@ -260,6 +293,12 @@ func (s *handlerManager) Trigger(handler Invocable) error { s.invokeCounter.WithLabelValues(handlerName, "rejected").Inc() logger.Warn("Dispatch queue is full, rejecting invocation", zap.Int("depth", len(queue))) + s.EmitFailure(message.ToDispatchInvoke(), &pb.InvocationFailure{ + Reason: pb.FailureReason_NO_RESOURCES, + Source: pb.FailureSource_RUNTIME, + Message: fmt.Sprintf("dispatch queue is full (%d waiting)", len(queue)), + CallerRejected: true, + }) return ErrDispatchQueueFull } s.queueDepthGauge.WithLabelValues(handlerName).Set(float64(len(queue))) @@ -268,6 +307,15 @@ func (s *handlerManager) Trigger(handler Invocable) error { defer cancel() <-message.Done() result, err := message.GetResult() + if errors.Is(err, context.DeadlineExceeded) && message.ClaimExpired() { + // No client picked it up in time, so it will never run. + s.EmitFailure(message.ToDispatchInvoke(), &pb.InvocationFailure{ + Reason: pb.FailureReason_NO_RESOURCES, + Source: pb.FailureSource_RUNTIME, + Message: fmt.Sprintf("expired after %v waiting for a free handler", + message.GetEntry().Timeout()), + }) + } s.queueDepthGauge.WithLabelValues(handlerName).Set(float64(len(queue))) s.handlerLatencyGauge.WithLabelValues(handlerName, result).Observe(float64(time.Since(startTime).Milliseconds())) if err != nil { @@ -282,6 +330,82 @@ func (s *handlerManager) Trigger(handler Invocable) error { return nil } +func (s *handlerManager) SubscribeFailures(dispatchId string) { + s.failureQueue(dispatchId, true) +} + +func (s *handlerManager) failureQueue(dispatchId string, create bool) *failureQueue { + s.failureLock.Lock() + defer s.failureLock.Unlock() + q, ok := s.failureQueues[dispatchId] + if !ok && create { + q = newFailureQueue() + s.failureQueues[dispatchId] = q + } + return q +} + +func (s *handlerManager) EmitFailure(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) { + logger := s.logger.With( + zap.String("handler", invoke.HandlerName), + zap.String("invocation-id", invoke.InvocationId), + zap.String("reason", failure.Reason.String()), + zap.String("source", failure.Source.String()), + zap.String("message", failure.Message), + ) + q := s.failureQueue(invoke.DispatchId, false) + if q == nil { + logger.Warn("Invocation failed, client does not handle failure events") + return + } + logger.Warn("Invocation failed, queueing failure event") + q.push(&pb.DispatchEvent{Invoke: invoke, Failure: failure}) + s.failuresPending.Inc() +} + +func (s *handlerManager) NextFailure(ctx context.Context, dispatchId string) (*pb.DispatchEvent, error) { + ev, err := s.failureQueue(dispatchId, true).pop(ctx) + if err != nil { + return nil, err + } + s.failuresPending.Dec() + return ev, nil +} + +func (s *handlerManager) RequeueFailure(dispatchId string, ev *pb.DispatchEvent) { + s.failureQueue(dispatchId, true).pushFront(ev) + s.failuresPending.Inc() +} + +func (s *handlerManager) Retrigger(invoke *pb.DispatchHandlerInvoke) (string, error) { + attempt := max(invoke.Attempt, 1) + if attempt >= maxRetryAttempts { + return "", fmt.Errorf("%w: attempt %d of %d", ErrRetryLimit, attempt, maxRetryAttempts) + } + + entry := s.handlers[invoke.HandlerId] + if entry == nil { + // The client re-registered since the invocation was sent, which gives + // the handler a new id under the same dispatch id and name. + for _, e := range s.handlers { + if e.DispatchId() == invoke.DispatchId && e.Name() == invoke.HandlerName { + entry = e + break + } + } + } + if entry == nil { + return "", fmt.Errorf("handler not found: %s", invoke.HandlerName) + } + + retry := NewHandlerInvoke(entry, invoke.Reason, invoke.Args) + retry.Attempt = attempt + 1 + if err := s.Trigger(retry); err != nil { + return "", err + } + return retry.Id, nil +} + func (s *handlerManager) getDispatchQueue(DispatchId string) chan Invocable { queue, ok := s.dispatchQueues[DispatchId] diff --git a/agent/server/handler/scheduled_handler_entry_test.go b/agent/server/handler/scheduled_handler_entry_test.go index 5d43e462..9f71f25f 100644 --- a/agent/server/handler/scheduled_handler_entry_test.go +++ b/agent/server/handler/scheduled_handler_entry_test.go @@ -287,3 +287,23 @@ func (fc *fakeCron) Add(spec string, cmd func()) (string, error) { func (fc *fakeCron) Remove(id string) { delete(fc.handlers, id) } + +func (fhm *fakeManager) SubscribeFailures(dispatchId string) { + panic("not implemented") // TODO: Implement +} + +func (fhm *fakeManager) EmitFailure(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) { + panic("not implemented") // TODO: Implement +} + +func (fhm *fakeManager) NextFailure(ctx context.Context, dispatchId string) (*pb.DispatchEvent, error) { + panic("not implemented") // TODO: Implement +} + +func (fhm *fakeManager) RequeueFailure(dispatchId string, ev *pb.DispatchEvent) { + panic("not implemented") // TODO: Implement +} + +func (fhm *fakeManager) Retrigger(invoke *pb.DispatchHandlerInvoke) (string, error) { + panic("not implemented") // TODO: Implement +} diff --git a/agent/server/http/http_server_webhook_test.go b/agent/server/http/http_server_webhook_test.go index 2d34de04..28a2b560 100644 --- a/agent/server/http/http_server_webhook_test.go +++ b/agent/server/http/http_server_webhook_test.go @@ -89,6 +89,7 @@ func TestHandleWebhookRejectsWhenQueueIsFull(t *testing.T) { _, err := manager.RegisterHandler("1", "test", time.Minute, option) require.NoError(t, err) require.NoError(t, manager.Start("1")) + manager.SubscribeFailures("1") webhookHandler := NewWebhookHandler(config.AgentConfig{}, logger, manager, nil) router := mux.NewRouter() @@ -121,4 +122,14 @@ func TestHandleWebhookRejectsWhenQueueIsFull(t *testing.T) { case <-time.After(5 * time.Second): require.Fail(t, "webhook request hung on a full queue instead of returning 503") } + + // The refused webhook still reaches the client's failure hook, flagged so + // it can tell the sender was asked to retry. + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + ev, err := manager.NextFailure(ctx, "1") + require.NoError(t, err) + require.Equal(t, pb.FailureReason_NO_RESOURCES, ev.Failure.Reason) + require.True(t, ev.Failure.CallerRejected) + require.Equal(t, "payload", ev.Invoke.Args["body"]) } diff --git a/agent/server/server.go b/agent/server/server.go index 3c0d74eb..7f80f90b 100644 --- a/agent/server/server.go +++ b/agent/server/server.go @@ -45,6 +45,10 @@ type AxonAgent struct { historyManager handler.HistoryManager } +// reportGrace is how long past an invocation's timeout the agent waits for +// the client's report before deciding the client never answered. +var reportGrace = 5 * time.Second + type inflightRequest struct { invocable handler.Invocable sentAt time.Time @@ -133,6 +137,12 @@ func (s *AxonAgent) Dispatch(stream pb.AxonAgent_DispatchServer) error { return err } + // Subscribe before any handler starts, so a failure can't happen + // before the client's DispatchEvents stream connects and be lost. + if req.FailureEvents { + s.Manager.SubscribeFailures(req.DispatchId) + } + s.sendInvocations(req.DispatchId, stream) } @@ -187,6 +197,12 @@ func (s *AxonAgent) sendInvocations(dispatchId string, stream pb.AxonAgent_Dispa continue } + if !invoke.ClaimDispatch() { + // It expired while queued and its failure was already + // reported; sending it now would run stale work. + continue + } + now := time.Now() msg := invoke.ToDispatchInvoke() @@ -205,16 +221,25 @@ func (s *AxonAgent) sendInvocations(dispatchId string, stream pb.AxonAgent_Dispa sentAt: now, }) go func(requestId string) { - <-time.After(time.Duration(msg.TimeoutMs) * time.Millisecond) - s.ReportInvocation(context.Background(), &pb.ReportInvocationRequest{ + // The grace lets a client that enforces the timeout itself + // report it first, so the failure is described once, by the + // side that saw it. + <-time.After(time.Duration(msg.TimeoutMs)*time.Millisecond + reportGrace) + message := fmt.Sprintf("handler did not report a result within %dms", msg.TimeoutMs) + completed := s.completeInvocation(context.Background(), &pb.ReportInvocationRequest{ HandlerInvoke: msg, Message: &pb.ReportInvocationRequest_Error{Error: &pb.Error{ - Code: "timeout", - Message: fmt.Sprintf( - "handler did not report a result within %dms", - msg.TimeoutMs), + Code: "timeout", + Message: message, }}, }) + if completed { + s.Manager.EmitFailure(msg, &pb.InvocationFailure{ + Reason: pb.FailureReason_TIMEOUT, + Source: pb.FailureSource_RUNTIME, + Message: message, + }) + } }(msg.InvocationId) } }() @@ -252,14 +277,21 @@ func describeError(e *pb.Error) string { // ReportInvocation is called by the client to report the result of an invocation, which will // log the result of an invocation into the history path. func (s *AxonAgent) ReportInvocation(ctx context.Context, req *pb.ReportInvocationRequest) (*pb.ReportInvocationResponse, error) { - s.inflightLock.RLock() + s.completeInvocation(ctx, req) + return &pb.ReportInvocationResponse{}, nil +} + +// completeInvocation records the result of an outstanding invocation. It +// returns false when the invocation was already completed by an earlier report. +func (s *AxonAgent) completeInvocation(ctx context.Context, req *pb.ReportInvocationRequest) bool { + s.inflightLock.Lock() ifr, ok := s.outstandingRequests[req.HandlerInvoke.InvocationId] - s.inflightLock.RUnlock() + delete(s.outstandingRequests, req.HandlerInvoke.InvocationId) + s.inflightLock.Unlock() if !ok { - return &pb.ReportInvocationResponse{}, nil + return false } - defer s.setOutstandingRequest(req.HandlerInvoke.InvocationId, nil) var requestErr error @@ -278,7 +310,50 @@ func (s *AxonAgent) ReportInvocation(ctx context.Context, req *pb.ReportInvocati if err != nil { s.logger.Error("failed to write history file", zap.Error(err)) } - return &pb.ReportInvocationResponse{}, nil + return true +} + +// DispatchEvents streams the agent's failure events for a dispatch id. An +// event whose send fails goes back on the queue for the next stream. +func (s *AxonAgent) DispatchEvents(req *pb.DispatchRequest, stream pb.AxonAgent_DispatchEventsServer) error { + if s.Manager == nil { + return fmt.Errorf("handler manager is not initialized") + } + + s.logger.Info("client subscribed to failure events", zap.String("dispatch-id", req.DispatchId)) + s.Manager.SubscribeFailures(req.DispatchId) + for { + ev, err := s.Manager.NextFailure(stream.Context(), req.DispatchId) + if err != nil { + return nil + } + if err := stream.Send(ev); err != nil { + s.Manager.RequeueFailure(req.DispatchId, ev) + s.logger.Warn("failed to send failure event, requeued", + zap.String("dispatch-id", req.DispatchId), zap.Error(err)) + return err + } + } +} + +// RetryInvocation queues a failed invocation again as its next attempt. +func (s *AxonAgent) RetryInvocation(ctx context.Context, req *pb.RetryInvocationRequest) (*pb.RetryInvocationResponse, error) { + if s.Manager == nil { + return nil, fmt.Errorf("handler manager is not initialized") + } + if req.Invoke == nil { + return &pb.RetryInvocationResponse{Error: &pb.Error{Code: "invalid", Message: "invoke is required"}}, nil + } + + id, err := s.Manager.Retrigger(req.Invoke) + if err != nil { + s.logger.Warn("retry rejected", + zap.String("handler", req.Invoke.HandlerName), + zap.String("invocation-id", req.Invoke.InvocationId), + zap.Error(err)) + return &pb.RetryInvocationResponse{Error: &pb.Error{Code: "rejected", Message: err.Error()}}, nil + } + return &pb.RetryInvocationResponse{InvocationId: id}, nil } // ListHandlers returns a list of all registered handlers diff --git a/agent/server/test_axon_client.go b/agent/server/test_axon_client.go index 362a07b9..d894abb0 100644 --- a/agent/server/test_axon_client.go +++ b/agent/server/test_axon_client.go @@ -21,6 +21,8 @@ type clientResult struct { done bool err error result string + // noReport skips ReportInvocation, like a client that hung. + noReport bool } type testAxonClient struct { @@ -30,6 +32,8 @@ type testAxonClient struct { callback clientCallback finished bool t *testing.T + + failureEvents bool } func NewTestAxonClient(t *testing.T, port int32, callback clientCallback) *testAxonClient { @@ -76,7 +80,8 @@ func (c *testAxonClient) Run(ctx context.Context) error { } dispatchReq := &pb.DispatchRequest{ - DispatchId: c.dispatchId, + DispatchId: c.dispatchId, + FailureEvents: c.failureEvents, } logger.Info("sending dispatch request") @@ -103,7 +108,7 @@ func (c *testAxonClient) Run(ctx context.Context) error { end := time.Now() invoke := m.GetInvoke() - if invoke != nil { + if invoke != nil && !result.noReport { report := &pb.ReportInvocationRequest{ HandlerInvoke: invoke, @@ -126,6 +131,10 @@ func (c *testAxonClient) Run(ctx context.Context) error { } _, err = c.client.ReportInvocation(ctx, report) + if ctx.Err() != nil { + // The test finished while this report was in flight. + return nil + } require.NoError(c.t, err) } } @@ -152,3 +161,11 @@ func (c *testAxonClient) GetHandlerHistory(ctx context.Context, name string) ([] } return resp.History, nil } + +func (c *testAxonClient) DispatchEvents(ctx context.Context) (grpc.ServerStreamingClient[pb.DispatchEvent], error) { + return c.client.DispatchEvents(ctx, &pb.DispatchRequest{DispatchId: c.dispatchId, FailureEvents: true}) +} + +func (c *testAxonClient) RetryInvocation(ctx context.Context, invoke *pb.DispatchHandlerInvoke) (*pb.RetryInvocationResponse, error) { + return c.client.RetryInvocation(ctx, &pb.RetryInvocationRequest{Invoke: invoke}) +} diff --git a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go index ea8b32af..5888e0df 100644 --- a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go +++ b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go @@ -1,7 +1,7 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: // protoc-gen-go v1.36.6 -// protoc v3.21.12 +// protoc v7.35.0 // source: cortex-axon-agent.proto package axon @@ -126,6 +126,113 @@ func (DispatchMessageType) EnumDescriptor() ([]byte, []int) { return file_cortex_axon_agent_proto_rawDescGZIP(), []int{1} } +type FailureReason int32 + +const ( + FailureReason_FAILURE_REASON_UNSPECIFIED FailureReason = 0 + // The handler returned an error or raised. + FailureReason_HANDLER_ERROR FailureReason = 1 + // The invocation did not finish within its timeout. + FailureReason_TIMEOUT FailureReason = 2 + // The invocation could not be dispatched: the queue was full, or it expired + // while waiting for a free handler. + FailureReason_NO_RESOURCES FailureReason = 3 +) + +// Enum value maps for FailureReason. +var ( + FailureReason_name = map[int32]string{ + 0: "FAILURE_REASON_UNSPECIFIED", + 1: "HANDLER_ERROR", + 2: "TIMEOUT", + 3: "NO_RESOURCES", + } + FailureReason_value = map[string]int32{ + "FAILURE_REASON_UNSPECIFIED": 0, + "HANDLER_ERROR": 1, + "TIMEOUT": 2, + "NO_RESOURCES": 3, + } +) + +func (x FailureReason) Enum() *FailureReason { + p := new(FailureReason) + *p = x + return p +} + +func (x FailureReason) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (FailureReason) Descriptor() protoreflect.EnumDescriptor { + return file_cortex_axon_agent_proto_enumTypes[2].Descriptor() +} + +func (FailureReason) Type() protoreflect.EnumType { + return &file_cortex_axon_agent_proto_enumTypes[2] +} + +func (x FailureReason) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use FailureReason.Descriptor instead. +func (FailureReason) EnumDescriptor() ([]byte, []int) { + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{2} +} + +type FailureSource int32 + +const ( + FailureSource_FAILURE_SOURCE_UNSPECIFIED FailureSource = 0 + // The failure came from handler code. + FailureSource_APPLICATION FailureSource = 1 + // The failure came from Axon: the agent, the SDK, or a Cortex API call. + FailureSource_RUNTIME FailureSource = 2 +) + +// Enum value maps for FailureSource. +var ( + FailureSource_name = map[int32]string{ + 0: "FAILURE_SOURCE_UNSPECIFIED", + 1: "APPLICATION", + 2: "RUNTIME", + } + FailureSource_value = map[string]int32{ + "FAILURE_SOURCE_UNSPECIFIED": 0, + "APPLICATION": 1, + "RUNTIME": 2, + } +) + +func (x FailureSource) Enum() *FailureSource { + p := new(FailureSource) + *p = x + return p +} + +func (x FailureSource) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (FailureSource) Descriptor() protoreflect.EnumDescriptor { + return file_cortex_axon_agent_proto_enumTypes[3].Descriptor() +} + +func (FailureSource) Type() protoreflect.EnumType { + return &file_cortex_axon_agent_proto_enumTypes[3] +} + +func (x FailureSource) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use FailureSource.Descriptor instead. +func (FailureSource) EnumDescriptor() ([]byte, []int) { + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{3} +} + type RegisterHandlerRequest struct { state protoimpl.MessageState `protogen:"open.v1"` DispatchId string `protobuf:"bytes,1,opt,name=dispatch_id,json=dispatchId,proto3" json:"dispatch_id,omitempty"` @@ -628,6 +735,9 @@ type DispatchRequest struct { state protoimpl.MessageState `protogen:"open.v1"` DispatchId string `protobuf:"bytes,1,opt,name=dispatch_id,json=dispatchId,proto3" json:"dispatch_id,omitempty"` ClientVersion string `protobuf:"bytes,2,opt,name=client_version,json=clientVersion,proto3" json:"client_version,omitempty"` + // The client handles failure events. The agent only keeps them for clients + // that ask, so an older client never has events pile up for it. + FailureEvents bool `protobuf:"varint,3,opt,name=failure_events,json=failureEvents,proto3" json:"failure_events,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -676,6 +786,13 @@ func (x *DispatchRequest) GetClientVersion() string { return "" } +func (x *DispatchRequest) GetFailureEvents() bool { + if x != nil { + return x.FailureEvents + } + return false +} + type DispatchMessage struct { state protoimpl.MessageState `protogen:"open.v1"` Type DispatchMessageType `protobuf:"varint,1,opt,name=type,proto3,enum=cortex.axon.DispatchMessageType" json:"type,omitempty"` @@ -751,14 +868,16 @@ type DispatchMessage_Invoke struct { func (*DispatchMessage_Invoke) isDispatchMessage_Message() {} type DispatchHandlerInvoke struct { - state protoimpl.MessageState `protogen:"open.v1"` - InvocationId string `protobuf:"bytes,1,opt,name=invocation_id,json=invocationId,proto3" json:"invocation_id,omitempty"` - DispatchId string `protobuf:"bytes,2,opt,name=dispatch_id,json=dispatchId,proto3" json:"dispatch_id,omitempty"` - HandlerId string `protobuf:"bytes,3,opt,name=handler_id,json=handlerId,proto3" json:"handler_id,omitempty"` - HandlerName string `protobuf:"bytes,4,opt,name=handler_name,json=handlerName,proto3" json:"handler_name,omitempty"` - TimeoutMs int32 `protobuf:"varint,10,opt,name=timeout_ms,json=timeoutMs,proto3" json:"timeout_ms,omitempty"` - Reason HandlerInvokeType `protobuf:"varint,11,opt,name=reason,proto3,enum=cortex.axon.HandlerInvokeType" json:"reason,omitempty"` - Args map[string]string `protobuf:"bytes,20,rep,name=args,proto3" json:"args,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + state protoimpl.MessageState `protogen:"open.v1"` + InvocationId string `protobuf:"bytes,1,opt,name=invocation_id,json=invocationId,proto3" json:"invocation_id,omitempty"` + DispatchId string `protobuf:"bytes,2,opt,name=dispatch_id,json=dispatchId,proto3" json:"dispatch_id,omitempty"` + HandlerId string `protobuf:"bytes,3,opt,name=handler_id,json=handlerId,proto3" json:"handler_id,omitempty"` + HandlerName string `protobuf:"bytes,4,opt,name=handler_name,json=handlerName,proto3" json:"handler_name,omitempty"` + TimeoutMs int32 `protobuf:"varint,10,opt,name=timeout_ms,json=timeoutMs,proto3" json:"timeout_ms,omitempty"` + Reason HandlerInvokeType `protobuf:"varint,11,opt,name=reason,proto3,enum=cortex.axon.HandlerInvokeType" json:"reason,omitempty"` + Args map[string]string `protobuf:"bytes,20,rep,name=args,proto3" json:"args,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + // 1 for the first attempt; 0 means the agent didn't say, so treat it as 1. + Attempt int32 `protobuf:"varint,30,opt,name=attempt,proto3" json:"attempt,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -842,6 +961,231 @@ func (x *DispatchHandlerInvoke) GetArgs() map[string]string { return nil } +func (x *DispatchHandlerInvoke) GetAttempt() int32 { + if x != nil { + return x.Attempt + } + return 0 +} + +type InvocationFailure struct { + state protoimpl.MessageState `protogen:"open.v1"` + Reason FailureReason `protobuf:"varint,1,opt,name=reason,proto3,enum=cortex.axon.FailureReason" json:"reason,omitempty"` + Source FailureSource `protobuf:"varint,2,opt,name=source,proto3,enum=cortex.axon.FailureSource" json:"source,omitempty"` + Message string `protobuf:"bytes,3,opt,name=message,proto3" json:"message,omitempty"` + // Whoever submitted the invocation was told it was not accepted (e.g. a + // webhook sender got a 503), so it may be sent again. + CallerRejected bool `protobuf:"varint,4,opt,name=caller_rejected,json=callerRejected,proto3" json:"caller_rejected,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *InvocationFailure) Reset() { + *x = InvocationFailure{} + mi := &file_cortex_axon_agent_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *InvocationFailure) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*InvocationFailure) ProtoMessage() {} + +func (x *InvocationFailure) ProtoReflect() protoreflect.Message { + mi := &file_cortex_axon_agent_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use InvocationFailure.ProtoReflect.Descriptor instead. +func (*InvocationFailure) Descriptor() ([]byte, []int) { + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{12} +} + +func (x *InvocationFailure) GetReason() FailureReason { + if x != nil { + return x.Reason + } + return FailureReason_FAILURE_REASON_UNSPECIFIED +} + +func (x *InvocationFailure) GetSource() FailureSource { + if x != nil { + return x.Source + } + return FailureSource_FAILURE_SOURCE_UNSPECIFIED +} + +func (x *InvocationFailure) GetMessage() string { + if x != nil { + return x.Message + } + return "" +} + +func (x *InvocationFailure) GetCallerRejected() bool { + if x != nil { + return x.CallerRejected + } + return false +} + +type DispatchEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Invoke *DispatchHandlerInvoke `protobuf:"bytes,1,opt,name=invoke,proto3" json:"invoke,omitempty"` + Failure *InvocationFailure `protobuf:"bytes,2,opt,name=failure,proto3" json:"failure,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DispatchEvent) Reset() { + *x = DispatchEvent{} + mi := &file_cortex_axon_agent_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DispatchEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DispatchEvent) ProtoMessage() {} + +func (x *DispatchEvent) ProtoReflect() protoreflect.Message { + mi := &file_cortex_axon_agent_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DispatchEvent.ProtoReflect.Descriptor instead. +func (*DispatchEvent) Descriptor() ([]byte, []int) { + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{13} +} + +func (x *DispatchEvent) GetInvoke() *DispatchHandlerInvoke { + if x != nil { + return x.Invoke + } + return nil +} + +func (x *DispatchEvent) GetFailure() *InvocationFailure { + if x != nil { + return x.Failure + } + return nil +} + +type RetryInvocationRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Invoke *DispatchHandlerInvoke `protobuf:"bytes,1,opt,name=invoke,proto3" json:"invoke,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RetryInvocationRequest) Reset() { + *x = RetryInvocationRequest{} + mi := &file_cortex_axon_agent_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RetryInvocationRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RetryInvocationRequest) ProtoMessage() {} + +func (x *RetryInvocationRequest) ProtoReflect() protoreflect.Message { + mi := &file_cortex_axon_agent_proto_msgTypes[14] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RetryInvocationRequest.ProtoReflect.Descriptor instead. +func (*RetryInvocationRequest) Descriptor() ([]byte, []int) { + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{14} +} + +func (x *RetryInvocationRequest) GetInvoke() *DispatchHandlerInvoke { + if x != nil { + return x.Invoke + } + return nil +} + +type RetryInvocationResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Error *Error `protobuf:"bytes,1,opt,name=error,proto3" json:"error,omitempty"` + InvocationId string `protobuf:"bytes,2,opt,name=invocation_id,json=invocationId,proto3" json:"invocation_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RetryInvocationResponse) Reset() { + *x = RetryInvocationResponse{} + mi := &file_cortex_axon_agent_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RetryInvocationResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RetryInvocationResponse) ProtoMessage() {} + +func (x *RetryInvocationResponse) ProtoReflect() protoreflect.Message { + mi := &file_cortex_axon_agent_proto_msgTypes[15] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RetryInvocationResponse.ProtoReflect.Descriptor instead. +func (*RetryInvocationResponse) Descriptor() ([]byte, []int) { + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{15} +} + +func (x *RetryInvocationResponse) GetError() *Error { + if x != nil { + return x.Error + } + return nil +} + +func (x *RetryInvocationResponse) GetInvocationId() string { + if x != nil { + return x.InvocationId + } + return "" +} + type Log struct { state protoimpl.MessageState `protogen:"open.v1"` Level string `protobuf:"bytes,1,opt,name=level,proto3" json:"level,omitempty"` @@ -853,7 +1197,7 @@ type Log struct { func (x *Log) Reset() { *x = Log{} - mi := &file_cortex_axon_agent_proto_msgTypes[12] + mi := &file_cortex_axon_agent_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -865,7 +1209,7 @@ func (x *Log) String() string { func (*Log) ProtoMessage() {} func (x *Log) ProtoReflect() protoreflect.Message { - mi := &file_cortex_axon_agent_proto_msgTypes[12] + mi := &file_cortex_axon_agent_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -878,7 +1222,7 @@ func (x *Log) ProtoReflect() protoreflect.Message { // Deprecated: Use Log.ProtoReflect.Descriptor instead. func (*Log) Descriptor() ([]byte, []int) { - return file_cortex_axon_agent_proto_rawDescGZIP(), []int{12} + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{16} } func (x *Log) GetLevel() string { @@ -919,7 +1263,7 @@ type ReportInvocationRequest struct { func (x *ReportInvocationRequest) Reset() { *x = ReportInvocationRequest{} - mi := &file_cortex_axon_agent_proto_msgTypes[13] + mi := &file_cortex_axon_agent_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -931,7 +1275,7 @@ func (x *ReportInvocationRequest) String() string { func (*ReportInvocationRequest) ProtoMessage() {} func (x *ReportInvocationRequest) ProtoReflect() protoreflect.Message { - mi := &file_cortex_axon_agent_proto_msgTypes[13] + mi := &file_cortex_axon_agent_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -944,7 +1288,7 @@ func (x *ReportInvocationRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ReportInvocationRequest.ProtoReflect.Descriptor instead. func (*ReportInvocationRequest) Descriptor() ([]byte, []int) { - return file_cortex_axon_agent_proto_rawDescGZIP(), []int{13} + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{17} } func (x *ReportInvocationRequest) GetHandlerInvoke() *DispatchHandlerInvoke { @@ -1025,7 +1369,7 @@ type ReportInvocationResponse struct { func (x *ReportInvocationResponse) Reset() { *x = ReportInvocationResponse{} - mi := &file_cortex_axon_agent_proto_msgTypes[14] + mi := &file_cortex_axon_agent_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1037,7 +1381,7 @@ func (x *ReportInvocationResponse) String() string { func (*ReportInvocationResponse) ProtoMessage() {} func (x *ReportInvocationResponse) ProtoReflect() protoreflect.Message { - mi := &file_cortex_axon_agent_proto_msgTypes[14] + mi := &file_cortex_axon_agent_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1050,7 +1394,7 @@ func (x *ReportInvocationResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ReportInvocationResponse.ProtoReflect.Descriptor instead. func (*ReportInvocationResponse) Descriptor() ([]byte, []int) { - return file_cortex_axon_agent_proto_rawDescGZIP(), []int{14} + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{18} } func (x *ReportInvocationResponse) GetError() *Error { @@ -1069,7 +1413,7 @@ type InvokeResult struct { func (x *InvokeResult) Reset() { *x = InvokeResult{} - mi := &file_cortex_axon_agent_proto_msgTypes[15] + mi := &file_cortex_axon_agent_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1081,7 +1425,7 @@ func (x *InvokeResult) String() string { func (*InvokeResult) ProtoMessage() {} func (x *InvokeResult) ProtoReflect() protoreflect.Message { - mi := &file_cortex_axon_agent_proto_msgTypes[15] + mi := &file_cortex_axon_agent_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1094,7 +1438,7 @@ func (x *InvokeResult) ProtoReflect() protoreflect.Message { // Deprecated: Use InvokeResult.ProtoReflect.Descriptor instead. func (*InvokeResult) Descriptor() ([]byte, []int) { - return file_cortex_axon_agent_proto_rawDescGZIP(), []int{15} + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{19} } func (x *InvokeResult) GetValue() string { @@ -1117,7 +1461,7 @@ type GetHandlerHistoryRequest struct { func (x *GetHandlerHistoryRequest) Reset() { *x = GetHandlerHistoryRequest{} - mi := &file_cortex_axon_agent_proto_msgTypes[16] + mi := &file_cortex_axon_agent_proto_msgTypes[20] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1129,7 +1473,7 @@ func (x *GetHandlerHistoryRequest) String() string { func (*GetHandlerHistoryRequest) ProtoMessage() {} func (x *GetHandlerHistoryRequest) ProtoReflect() protoreflect.Message { - mi := &file_cortex_axon_agent_proto_msgTypes[16] + mi := &file_cortex_axon_agent_proto_msgTypes[20] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1142,7 +1486,7 @@ func (x *GetHandlerHistoryRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetHandlerHistoryRequest.ProtoReflect.Descriptor instead. func (*GetHandlerHistoryRequest) Descriptor() ([]byte, []int) { - return file_cortex_axon_agent_proto_rawDescGZIP(), []int{16} + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{20} } func (x *GetHandlerHistoryRequest) GetHandlerName() string { @@ -1198,7 +1542,7 @@ type HandlerExecution struct { func (x *HandlerExecution) Reset() { *x = HandlerExecution{} - mi := &file_cortex_axon_agent_proto_msgTypes[17] + mi := &file_cortex_axon_agent_proto_msgTypes[21] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1210,7 +1554,7 @@ func (x *HandlerExecution) String() string { func (*HandlerExecution) ProtoMessage() {} func (x *HandlerExecution) ProtoReflect() protoreflect.Message { - mi := &file_cortex_axon_agent_proto_msgTypes[17] + mi := &file_cortex_axon_agent_proto_msgTypes[21] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1223,7 +1567,7 @@ func (x *HandlerExecution) ProtoReflect() protoreflect.Message { // Deprecated: Use HandlerExecution.ProtoReflect.Descriptor instead. func (*HandlerExecution) Descriptor() ([]byte, []int) { - return file_cortex_axon_agent_proto_rawDescGZIP(), []int{17} + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{21} } func (x *HandlerExecution) GetHandlerName() string { @@ -1306,7 +1650,7 @@ type GetHandlerHistoryResponse struct { func (x *GetHandlerHistoryResponse) Reset() { *x = GetHandlerHistoryResponse{} - mi := &file_cortex_axon_agent_proto_msgTypes[18] + mi := &file_cortex_axon_agent_proto_msgTypes[22] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1318,7 +1662,7 @@ func (x *GetHandlerHistoryResponse) String() string { func (*GetHandlerHistoryResponse) ProtoMessage() {} func (x *GetHandlerHistoryResponse) ProtoReflect() protoreflect.Message { - mi := &file_cortex_axon_agent_proto_msgTypes[18] + mi := &file_cortex_axon_agent_proto_msgTypes[22] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1331,7 +1675,7 @@ func (x *GetHandlerHistoryResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetHandlerHistoryResponse.ProtoReflect.Descriptor instead. func (*GetHandlerHistoryResponse) Descriptor() ([]byte, []int) { - return file_cortex_axon_agent_proto_rawDescGZIP(), []int{18} + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{22} } func (x *GetHandlerHistoryResponse) GetError() *Error { @@ -1384,16 +1728,17 @@ const file_cortex_axon_agent_proto_rawDesc = "" + "\tis_active\x18e \x01(\bR\bisActive\"v\n" + "\x14ListHandlersResponse\x12(\n" + "\x05error\x18\x01 \x01(\v2\x12.cortex.axon.ErrorR\x05error\x124\n" + - "\bhandlers\x18\x02 \x03(\v2\x18.cortex.axon.HandlerInfoR\bhandlers\"Y\n" + + "\bhandlers\x18\x02 \x03(\v2\x18.cortex.axon.HandlerInfoR\bhandlers\"\x80\x01\n" + "\x0fDispatchRequest\x12\x1f\n" + "\vdispatch_id\x18\x01 \x01(\tR\n" + "dispatchId\x12%\n" + - "\x0eclient_version\x18\x02 \x01(\tR\rclientVersion\"\x90\x01\n" + + "\x0eclient_version\x18\x02 \x01(\tR\rclientVersion\x12%\n" + + "\x0efailure_events\x18\x03 \x01(\bR\rfailureEvents\"\x90\x01\n" + "\x0fDispatchMessage\x124\n" + "\x04type\x18\x01 \x01(\x0e2 .cortex.axon.DispatchMessageTypeR\x04type\x12<\n" + "\x06invoke\x18\n" + " \x01(\v2\".cortex.axon.DispatchHandlerInvokeH\x00R\x06invokeB\t\n" + - "\amessage\"\xf1\x02\n" + + "\amessage\"\x8b\x03\n" + "\x15DispatchHandlerInvoke\x12#\n" + "\rinvocation_id\x18\x01 \x01(\tR\finvocationId\x12\x1f\n" + "\vdispatch_id\x18\x02 \x01(\tR\n" + @@ -1405,10 +1750,24 @@ const file_cortex_axon_agent_proto_rawDesc = "" + "timeout_ms\x18\n" + " \x01(\x05R\ttimeoutMs\x126\n" + "\x06reason\x18\v \x01(\x0e2\x1e.cortex.axon.HandlerInvokeTypeR\x06reason\x12@\n" + - "\x04args\x18\x14 \x03(\v2,.cortex.axon.DispatchHandlerInvoke.ArgsEntryR\x04args\x1a7\n" + + "\x04args\x18\x14 \x03(\v2,.cortex.axon.DispatchHandlerInvoke.ArgsEntryR\x04args\x12\x18\n" + + "\aattempt\x18\x1e \x01(\x05R\aattempt\x1a7\n" + "\tArgsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + - "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"o\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xbe\x01\n" + + "\x11InvocationFailure\x122\n" + + "\x06reason\x18\x01 \x01(\x0e2\x1a.cortex.axon.FailureReasonR\x06reason\x122\n" + + "\x06source\x18\x02 \x01(\x0e2\x1a.cortex.axon.FailureSourceR\x06source\x12\x18\n" + + "\amessage\x18\x03 \x01(\tR\amessage\x12'\n" + + "\x0fcaller_rejected\x18\x04 \x01(\bR\x0ecallerRejected\"\x85\x01\n" + + "\rDispatchEvent\x12:\n" + + "\x06invoke\x18\x01 \x01(\v2\".cortex.axon.DispatchHandlerInvokeR\x06invoke\x128\n" + + "\afailure\x18\x02 \x01(\v2\x1e.cortex.axon.InvocationFailureR\afailure\"T\n" + + "\x16RetryInvocationRequest\x12:\n" + + "\x06invoke\x18\x01 \x01(\v2\".cortex.axon.DispatchHandlerInvokeR\x06invoke\"h\n" + + "\x17RetryInvocationResponse\x12(\n" + + "\x05error\x18\x01 \x01(\v2\x12.cortex.axon.ErrorR\x05error\x12#\n" + + "\rinvocation_id\x18\x02 \x01(\tR\finvocationId\"o\n" + "\x03Log\x12\x14\n" + "\x05level\x18\x01 \x01(\tR\x05level\x128\n" + "\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12\x18\n" + @@ -1462,14 +1821,25 @@ const file_cortex_axon_agent_proto_rawDesc = "" + "\x13DispatchMessageType\x12\x1e\n" + "\x1aDISPATCH_MESSAGE_TYPE_NONE\x10\x00\x12\x1b\n" + "\x17DISPATCH_MESSAGE_INVOKE\x10\x01\x12#\n" + - "\x1fDISPATCH_MESSAGE_WORK_COMPLETED\x10\x022\xb3\x04\n" + + "\x1fDISPATCH_MESSAGE_WORK_COMPLETED\x10\x02*a\n" + + "\rFailureReason\x12\x1e\n" + + "\x1aFAILURE_REASON_UNSPECIFIED\x10\x00\x12\x11\n" + + "\rHANDLER_ERROR\x10\x01\x12\v\n" + + "\aTIMEOUT\x10\x02\x12\x10\n" + + "\fNO_RESOURCES\x10\x03*M\n" + + "\rFailureSource\x12\x1e\n" + + "\x1aFAILURE_SOURCE_UNSPECIFIED\x10\x00\x12\x0f\n" + + "\vAPPLICATION\x10\x01\x12\v\n" + + "\aRUNTIME\x10\x022\xdf\x05\n" + "\tAxonAgent\x12\\\n" + "\x0fRegisterHandler\x12#.cortex.axon.RegisterHandlerRequest\x1a$.cortex.axon.RegisterHandlerResponse\x12b\n" + "\x11UnregisterHandler\x12%.cortex.axon.UnregisterHandlerRequest\x1a&.cortex.axon.UnregisterHandlerResponse\x12S\n" + "\fListHandlers\x12 .cortex.axon.ListHandlersRequest\x1a!.cortex.axon.ListHandlersResponse\x12b\n" + "\x11GetHandlerHistory\x12%.cortex.axon.GetHandlerHistoryRequest\x1a&.cortex.axon.GetHandlerHistoryResponse\x12J\n" + "\bDispatch\x12\x1c.cortex.axon.DispatchRequest\x1a\x1c.cortex.axon.DispatchMessage(\x010\x01\x12_\n" + - "\x10ReportInvocation\x12$.cortex.axon.ReportInvocationRequest\x1a%.cortex.axon.ReportInvocationResponseB\x1cZ\x1agithub.com/cortexapps/axonb\x06proto3" + "\x10ReportInvocation\x12$.cortex.axon.ReportInvocationRequest\x1a%.cortex.axon.ReportInvocationResponse\x12L\n" + + "\x0eDispatchEvents\x12\x1c.cortex.axon.DispatchRequest\x1a\x1a.cortex.axon.DispatchEvent0\x01\x12\\\n" + + "\x0fRetryInvocation\x12#.cortex.axon.RetryInvocationRequest\x1a$.cortex.axon.RetryInvocationResponseB\x1cZ\x1agithub.com/cortexapps/axonb\x06proto3" var ( file_cortex_axon_agent_proto_rawDescOnce sync.Once @@ -1483,81 +1853,97 @@ func file_cortex_axon_agent_proto_rawDescGZIP() []byte { return file_cortex_axon_agent_proto_rawDescData } -var file_cortex_axon_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 2) -var file_cortex_axon_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 20) +var file_cortex_axon_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 4) +var file_cortex_axon_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 24) var file_cortex_axon_agent_proto_goTypes = []any{ (HandlerInvokeType)(0), // 0: cortex.axon.HandlerInvokeType (DispatchMessageType)(0), // 1: cortex.axon.DispatchMessageType - (*RegisterHandlerRequest)(nil), // 2: cortex.axon.RegisterHandlerRequest - (*HandlerInvokeOption)(nil), // 3: cortex.axon.HandlerInvokeOption - (*HandlerOption)(nil), // 4: cortex.axon.HandlerOption - (*RegisterHandlerResponse)(nil), // 5: cortex.axon.RegisterHandlerResponse - (*UnregisterHandlerRequest)(nil), // 6: cortex.axon.UnregisterHandlerRequest - (*UnregisterHandlerResponse)(nil), // 7: cortex.axon.UnregisterHandlerResponse - (*ListHandlersRequest)(nil), // 8: cortex.axon.ListHandlersRequest - (*HandlerInfo)(nil), // 9: cortex.axon.HandlerInfo - (*ListHandlersResponse)(nil), // 10: cortex.axon.ListHandlersResponse - (*DispatchRequest)(nil), // 11: cortex.axon.DispatchRequest - (*DispatchMessage)(nil), // 12: cortex.axon.DispatchMessage - (*DispatchHandlerInvoke)(nil), // 13: cortex.axon.DispatchHandlerInvoke - (*Log)(nil), // 14: cortex.axon.Log - (*ReportInvocationRequest)(nil), // 15: cortex.axon.ReportInvocationRequest - (*ReportInvocationResponse)(nil), // 16: cortex.axon.ReportInvocationResponse - (*InvokeResult)(nil), // 17: cortex.axon.InvokeResult - (*GetHandlerHistoryRequest)(nil), // 18: cortex.axon.GetHandlerHistoryRequest - (*HandlerExecution)(nil), // 19: cortex.axon.HandlerExecution - (*GetHandlerHistoryResponse)(nil), // 20: cortex.axon.GetHandlerHistoryResponse - nil, // 21: cortex.axon.DispatchHandlerInvoke.ArgsEntry - (*Error)(nil), // 22: cortex.axon.Error - (*timestamppb.Timestamp)(nil), // 23: google.protobuf.Timestamp + (FailureReason)(0), // 2: cortex.axon.FailureReason + (FailureSource)(0), // 3: cortex.axon.FailureSource + (*RegisterHandlerRequest)(nil), // 4: cortex.axon.RegisterHandlerRequest + (*HandlerInvokeOption)(nil), // 5: cortex.axon.HandlerInvokeOption + (*HandlerOption)(nil), // 6: cortex.axon.HandlerOption + (*RegisterHandlerResponse)(nil), // 7: cortex.axon.RegisterHandlerResponse + (*UnregisterHandlerRequest)(nil), // 8: cortex.axon.UnregisterHandlerRequest + (*UnregisterHandlerResponse)(nil), // 9: cortex.axon.UnregisterHandlerResponse + (*ListHandlersRequest)(nil), // 10: cortex.axon.ListHandlersRequest + (*HandlerInfo)(nil), // 11: cortex.axon.HandlerInfo + (*ListHandlersResponse)(nil), // 12: cortex.axon.ListHandlersResponse + (*DispatchRequest)(nil), // 13: cortex.axon.DispatchRequest + (*DispatchMessage)(nil), // 14: cortex.axon.DispatchMessage + (*DispatchHandlerInvoke)(nil), // 15: cortex.axon.DispatchHandlerInvoke + (*InvocationFailure)(nil), // 16: cortex.axon.InvocationFailure + (*DispatchEvent)(nil), // 17: cortex.axon.DispatchEvent + (*RetryInvocationRequest)(nil), // 18: cortex.axon.RetryInvocationRequest + (*RetryInvocationResponse)(nil), // 19: cortex.axon.RetryInvocationResponse + (*Log)(nil), // 20: cortex.axon.Log + (*ReportInvocationRequest)(nil), // 21: cortex.axon.ReportInvocationRequest + (*ReportInvocationResponse)(nil), // 22: cortex.axon.ReportInvocationResponse + (*InvokeResult)(nil), // 23: cortex.axon.InvokeResult + (*GetHandlerHistoryRequest)(nil), // 24: cortex.axon.GetHandlerHistoryRequest + (*HandlerExecution)(nil), // 25: cortex.axon.HandlerExecution + (*GetHandlerHistoryResponse)(nil), // 26: cortex.axon.GetHandlerHistoryResponse + nil, // 27: cortex.axon.DispatchHandlerInvoke.ArgsEntry + (*Error)(nil), // 28: cortex.axon.Error + (*timestamppb.Timestamp)(nil), // 29: google.protobuf.Timestamp } var file_cortex_axon_agent_proto_depIdxs = []int32{ - 4, // 0: cortex.axon.RegisterHandlerRequest.options:type_name -> cortex.axon.HandlerOption + 6, // 0: cortex.axon.RegisterHandlerRequest.options:type_name -> cortex.axon.HandlerOption 0, // 1: cortex.axon.HandlerInvokeOption.type:type_name -> cortex.axon.HandlerInvokeType - 3, // 2: cortex.axon.HandlerOption.invoke:type_name -> cortex.axon.HandlerInvokeOption - 22, // 3: cortex.axon.RegisterHandlerResponse.error:type_name -> cortex.axon.Error - 22, // 4: cortex.axon.UnregisterHandlerResponse.error:type_name -> cortex.axon.Error - 4, // 5: cortex.axon.HandlerInfo.options:type_name -> cortex.axon.HandlerOption - 23, // 6: cortex.axon.HandlerInfo.last_invoked_client_timestamp:type_name -> google.protobuf.Timestamp - 22, // 7: cortex.axon.ListHandlersResponse.error:type_name -> cortex.axon.Error - 9, // 8: cortex.axon.ListHandlersResponse.handlers:type_name -> cortex.axon.HandlerInfo + 5, // 2: cortex.axon.HandlerOption.invoke:type_name -> cortex.axon.HandlerInvokeOption + 28, // 3: cortex.axon.RegisterHandlerResponse.error:type_name -> cortex.axon.Error + 28, // 4: cortex.axon.UnregisterHandlerResponse.error:type_name -> cortex.axon.Error + 6, // 5: cortex.axon.HandlerInfo.options:type_name -> cortex.axon.HandlerOption + 29, // 6: cortex.axon.HandlerInfo.last_invoked_client_timestamp:type_name -> google.protobuf.Timestamp + 28, // 7: cortex.axon.ListHandlersResponse.error:type_name -> cortex.axon.Error + 11, // 8: cortex.axon.ListHandlersResponse.handlers:type_name -> cortex.axon.HandlerInfo 1, // 9: cortex.axon.DispatchMessage.type:type_name -> cortex.axon.DispatchMessageType - 13, // 10: cortex.axon.DispatchMessage.invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 15, // 10: cortex.axon.DispatchMessage.invoke:type_name -> cortex.axon.DispatchHandlerInvoke 0, // 11: cortex.axon.DispatchHandlerInvoke.reason:type_name -> cortex.axon.HandlerInvokeType - 21, // 12: cortex.axon.DispatchHandlerInvoke.args:type_name -> cortex.axon.DispatchHandlerInvoke.ArgsEntry - 23, // 13: cortex.axon.Log.timestamp:type_name -> google.protobuf.Timestamp - 13, // 14: cortex.axon.ReportInvocationRequest.handler_invoke:type_name -> cortex.axon.DispatchHandlerInvoke - 23, // 15: cortex.axon.ReportInvocationRequest.start_client_timestamp:type_name -> google.protobuf.Timestamp - 17, // 16: cortex.axon.ReportInvocationRequest.result:type_name -> cortex.axon.InvokeResult - 22, // 17: cortex.axon.ReportInvocationRequest.error:type_name -> cortex.axon.Error - 14, // 18: cortex.axon.ReportInvocationRequest.logs:type_name -> cortex.axon.Log - 22, // 19: cortex.axon.ReportInvocationResponse.error:type_name -> cortex.axon.Error - 23, // 20: cortex.axon.GetHandlerHistoryRequest.start_time:type_name -> google.protobuf.Timestamp - 23, // 21: cortex.axon.GetHandlerHistoryRequest.end_time:type_name -> google.protobuf.Timestamp - 23, // 22: cortex.axon.HandlerExecution.publish_server_timestamp:type_name -> google.protobuf.Timestamp - 23, // 23: cortex.axon.HandlerExecution.receive_server_timestamp:type_name -> google.protobuf.Timestamp - 23, // 24: cortex.axon.HandlerExecution.start_client_timestamp:type_name -> google.protobuf.Timestamp - 22, // 25: cortex.axon.HandlerExecution.error:type_name -> cortex.axon.Error - 14, // 26: cortex.axon.HandlerExecution.logs:type_name -> cortex.axon.Log - 22, // 27: cortex.axon.GetHandlerHistoryResponse.error:type_name -> cortex.axon.Error - 19, // 28: cortex.axon.GetHandlerHistoryResponse.history:type_name -> cortex.axon.HandlerExecution - 2, // 29: cortex.axon.AxonAgent.RegisterHandler:input_type -> cortex.axon.RegisterHandlerRequest - 6, // 30: cortex.axon.AxonAgent.UnregisterHandler:input_type -> cortex.axon.UnregisterHandlerRequest - 8, // 31: cortex.axon.AxonAgent.ListHandlers:input_type -> cortex.axon.ListHandlersRequest - 18, // 32: cortex.axon.AxonAgent.GetHandlerHistory:input_type -> cortex.axon.GetHandlerHistoryRequest - 11, // 33: cortex.axon.AxonAgent.Dispatch:input_type -> cortex.axon.DispatchRequest - 15, // 34: cortex.axon.AxonAgent.ReportInvocation:input_type -> cortex.axon.ReportInvocationRequest - 5, // 35: cortex.axon.AxonAgent.RegisterHandler:output_type -> cortex.axon.RegisterHandlerResponse - 7, // 36: cortex.axon.AxonAgent.UnregisterHandler:output_type -> cortex.axon.UnregisterHandlerResponse - 10, // 37: cortex.axon.AxonAgent.ListHandlers:output_type -> cortex.axon.ListHandlersResponse - 20, // 38: cortex.axon.AxonAgent.GetHandlerHistory:output_type -> cortex.axon.GetHandlerHistoryResponse - 12, // 39: cortex.axon.AxonAgent.Dispatch:output_type -> cortex.axon.DispatchMessage - 16, // 40: cortex.axon.AxonAgent.ReportInvocation:output_type -> cortex.axon.ReportInvocationResponse - 35, // [35:41] is the sub-list for method output_type - 29, // [29:35] is the sub-list for method input_type - 29, // [29:29] is the sub-list for extension type_name - 29, // [29:29] is the sub-list for extension extendee - 0, // [0:29] is the sub-list for field type_name + 27, // 12: cortex.axon.DispatchHandlerInvoke.args:type_name -> cortex.axon.DispatchHandlerInvoke.ArgsEntry + 2, // 13: cortex.axon.InvocationFailure.reason:type_name -> cortex.axon.FailureReason + 3, // 14: cortex.axon.InvocationFailure.source:type_name -> cortex.axon.FailureSource + 15, // 15: cortex.axon.DispatchEvent.invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 16, // 16: cortex.axon.DispatchEvent.failure:type_name -> cortex.axon.InvocationFailure + 15, // 17: cortex.axon.RetryInvocationRequest.invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 28, // 18: cortex.axon.RetryInvocationResponse.error:type_name -> cortex.axon.Error + 29, // 19: cortex.axon.Log.timestamp:type_name -> google.protobuf.Timestamp + 15, // 20: cortex.axon.ReportInvocationRequest.handler_invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 29, // 21: cortex.axon.ReportInvocationRequest.start_client_timestamp:type_name -> google.protobuf.Timestamp + 23, // 22: cortex.axon.ReportInvocationRequest.result:type_name -> cortex.axon.InvokeResult + 28, // 23: cortex.axon.ReportInvocationRequest.error:type_name -> cortex.axon.Error + 20, // 24: cortex.axon.ReportInvocationRequest.logs:type_name -> cortex.axon.Log + 28, // 25: cortex.axon.ReportInvocationResponse.error:type_name -> cortex.axon.Error + 29, // 26: cortex.axon.GetHandlerHistoryRequest.start_time:type_name -> google.protobuf.Timestamp + 29, // 27: cortex.axon.GetHandlerHistoryRequest.end_time:type_name -> google.protobuf.Timestamp + 29, // 28: cortex.axon.HandlerExecution.publish_server_timestamp:type_name -> google.protobuf.Timestamp + 29, // 29: cortex.axon.HandlerExecution.receive_server_timestamp:type_name -> google.protobuf.Timestamp + 29, // 30: cortex.axon.HandlerExecution.start_client_timestamp:type_name -> google.protobuf.Timestamp + 28, // 31: cortex.axon.HandlerExecution.error:type_name -> cortex.axon.Error + 20, // 32: cortex.axon.HandlerExecution.logs:type_name -> cortex.axon.Log + 28, // 33: cortex.axon.GetHandlerHistoryResponse.error:type_name -> cortex.axon.Error + 25, // 34: cortex.axon.GetHandlerHistoryResponse.history:type_name -> cortex.axon.HandlerExecution + 4, // 35: cortex.axon.AxonAgent.RegisterHandler:input_type -> cortex.axon.RegisterHandlerRequest + 8, // 36: cortex.axon.AxonAgent.UnregisterHandler:input_type -> cortex.axon.UnregisterHandlerRequest + 10, // 37: cortex.axon.AxonAgent.ListHandlers:input_type -> cortex.axon.ListHandlersRequest + 24, // 38: cortex.axon.AxonAgent.GetHandlerHistory:input_type -> cortex.axon.GetHandlerHistoryRequest + 13, // 39: cortex.axon.AxonAgent.Dispatch:input_type -> cortex.axon.DispatchRequest + 21, // 40: cortex.axon.AxonAgent.ReportInvocation:input_type -> cortex.axon.ReportInvocationRequest + 13, // 41: cortex.axon.AxonAgent.DispatchEvents:input_type -> cortex.axon.DispatchRequest + 18, // 42: cortex.axon.AxonAgent.RetryInvocation:input_type -> cortex.axon.RetryInvocationRequest + 7, // 43: cortex.axon.AxonAgent.RegisterHandler:output_type -> cortex.axon.RegisterHandlerResponse + 9, // 44: cortex.axon.AxonAgent.UnregisterHandler:output_type -> cortex.axon.UnregisterHandlerResponse + 12, // 45: cortex.axon.AxonAgent.ListHandlers:output_type -> cortex.axon.ListHandlersResponse + 26, // 46: cortex.axon.AxonAgent.GetHandlerHistory:output_type -> cortex.axon.GetHandlerHistoryResponse + 14, // 47: cortex.axon.AxonAgent.Dispatch:output_type -> cortex.axon.DispatchMessage + 22, // 48: cortex.axon.AxonAgent.ReportInvocation:output_type -> cortex.axon.ReportInvocationResponse + 17, // 49: cortex.axon.AxonAgent.DispatchEvents:output_type -> cortex.axon.DispatchEvent + 19, // 50: cortex.axon.AxonAgent.RetryInvocation:output_type -> cortex.axon.RetryInvocationResponse + 43, // [43:51] is the sub-list for method output_type + 35, // [35:43] is the sub-list for method input_type + 35, // [35:35] is the sub-list for extension type_name + 35, // [35:35] is the sub-list for extension extendee + 0, // [0:35] is the sub-list for field type_name } func init() { file_cortex_axon_agent_proto_init() } @@ -1572,7 +1958,7 @@ func file_cortex_axon_agent_proto_init() { file_cortex_axon_agent_proto_msgTypes[10].OneofWrappers = []any{ (*DispatchMessage_Invoke)(nil), } - file_cortex_axon_agent_proto_msgTypes[13].OneofWrappers = []any{ + file_cortex_axon_agent_proto_msgTypes[17].OneofWrappers = []any{ (*ReportInvocationRequest_Result)(nil), (*ReportInvocationRequest_Error)(nil), } @@ -1581,8 +1967,8 @@ func file_cortex_axon_agent_proto_init() { File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_cortex_axon_agent_proto_rawDesc), len(file_cortex_axon_agent_proto_rawDesc)), - NumEnums: 2, - NumMessages: 20, + NumEnums: 4, + NumMessages: 24, NumExtensions: 0, NumServices: 1, }, diff --git a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go index 3d53df53..83f80606 100644 --- a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go +++ b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go @@ -1,7 +1,7 @@ // Code generated by protoc-gen-go-grpc. DO NOT EDIT. // versions: // - protoc-gen-go-grpc v1.5.1 -// - protoc v3.21.12 +// - protoc v7.35.0 // source: cortex-axon-agent.proto package axon @@ -25,6 +25,8 @@ const ( AxonAgent_GetHandlerHistory_FullMethodName = "/cortex.axon.AxonAgent/GetHandlerHistory" AxonAgent_Dispatch_FullMethodName = "/cortex.axon.AxonAgent/Dispatch" AxonAgent_ReportInvocation_FullMethodName = "/cortex.axon.AxonAgent/ReportInvocation" + AxonAgent_DispatchEvents_FullMethodName = "/cortex.axon.AxonAgent/DispatchEvents" + AxonAgent_RetryInvocation_FullMethodName = "/cortex.axon.AxonAgent/RetryInvocation" ) // AxonAgentClient is the client API for AxonAgent service. @@ -37,6 +39,13 @@ type AxonAgentClient interface { GetHandlerHistory(ctx context.Context, in *GetHandlerHistoryRequest, opts ...grpc.CallOption) (*GetHandlerHistoryResponse, error) Dispatch(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[DispatchRequest, DispatchMessage], error) ReportInvocation(ctx context.Context, in *ReportInvocationRequest, opts ...grpc.CallOption) (*ReportInvocationResponse, error) + // DispatchEvents streams failures the agent detected on its own: invocations + // that never reached the client, or that the client never reported. They use + // their own stream so a client that has stopped reading Dispatch because + // every handler is busy still hears about them. + DispatchEvents(ctx context.Context, in *DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[DispatchEvent], error) + // RetryInvocation queues a failed invocation again as a new attempt. + RetryInvocation(ctx context.Context, in *RetryInvocationRequest, opts ...grpc.CallOption) (*RetryInvocationResponse, error) } type axonAgentClient struct { @@ -110,6 +119,35 @@ func (c *axonAgentClient) ReportInvocation(ctx context.Context, in *ReportInvoca return out, nil } +func (c *axonAgentClient) DispatchEvents(ctx context.Context, in *DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[DispatchEvent], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &AxonAgent_ServiceDesc.Streams[1], AxonAgent_DispatchEvents_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[DispatchRequest, DispatchEvent]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type AxonAgent_DispatchEventsClient = grpc.ServerStreamingClient[DispatchEvent] + +func (c *axonAgentClient) RetryInvocation(ctx context.Context, in *RetryInvocationRequest, opts ...grpc.CallOption) (*RetryInvocationResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RetryInvocationResponse) + err := c.cc.Invoke(ctx, AxonAgent_RetryInvocation_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // AxonAgentServer is the server API for AxonAgent service. // All implementations must embed UnimplementedAxonAgentServer // for forward compatibility. @@ -120,6 +158,13 @@ type AxonAgentServer interface { GetHandlerHistory(context.Context, *GetHandlerHistoryRequest) (*GetHandlerHistoryResponse, error) Dispatch(grpc.BidiStreamingServer[DispatchRequest, DispatchMessage]) error ReportInvocation(context.Context, *ReportInvocationRequest) (*ReportInvocationResponse, error) + // DispatchEvents streams failures the agent detected on its own: invocations + // that never reached the client, or that the client never reported. They use + // their own stream so a client that has stopped reading Dispatch because + // every handler is busy still hears about them. + DispatchEvents(*DispatchRequest, grpc.ServerStreamingServer[DispatchEvent]) error + // RetryInvocation queues a failed invocation again as a new attempt. + RetryInvocation(context.Context, *RetryInvocationRequest) (*RetryInvocationResponse, error) mustEmbedUnimplementedAxonAgentServer() } @@ -148,6 +193,12 @@ func (UnimplementedAxonAgentServer) Dispatch(grpc.BidiStreamingServer[DispatchRe func (UnimplementedAxonAgentServer) ReportInvocation(context.Context, *ReportInvocationRequest) (*ReportInvocationResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method ReportInvocation not implemented") } +func (UnimplementedAxonAgentServer) DispatchEvents(*DispatchRequest, grpc.ServerStreamingServer[DispatchEvent]) error { + return status.Errorf(codes.Unimplemented, "method DispatchEvents not implemented") +} +func (UnimplementedAxonAgentServer) RetryInvocation(context.Context, *RetryInvocationRequest) (*RetryInvocationResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method RetryInvocation not implemented") +} func (UnimplementedAxonAgentServer) mustEmbedUnimplementedAxonAgentServer() {} func (UnimplementedAxonAgentServer) testEmbeddedByValue() {} @@ -266,6 +317,35 @@ func _AxonAgent_ReportInvocation_Handler(srv interface{}, ctx context.Context, d return interceptor(ctx, in, info, handler) } +func _AxonAgent_DispatchEvents_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(DispatchRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(AxonAgentServer).DispatchEvents(m, &grpc.GenericServerStream[DispatchRequest, DispatchEvent]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type AxonAgent_DispatchEventsServer = grpc.ServerStreamingServer[DispatchEvent] + +func _AxonAgent_RetryInvocation_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RetryInvocationRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AxonAgentServer).RetryInvocation(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AxonAgent_RetryInvocation_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AxonAgentServer).RetryInvocation(ctx, req.(*RetryInvocationRequest)) + } + return interceptor(ctx, in, info, handler) +} + // AxonAgent_ServiceDesc is the grpc.ServiceDesc for AxonAgent service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -293,6 +373,10 @@ var AxonAgent_ServiceDesc = grpc.ServiceDesc{ MethodName: "ReportInvocation", Handler: _AxonAgent_ReportInvocation_Handler, }, + { + MethodName: "RetryInvocation", + Handler: _AxonAgent_RetryInvocation_Handler, + }, }, Streams: []grpc.StreamDesc{ { @@ -301,6 +385,11 @@ var AxonAgent_ServiceDesc = grpc.ServiceDesc{ ServerStreams: true, ClientStreams: true, }, + { + StreamName: "DispatchEvents", + Handler: _AxonAgent_DispatchEvents_Handler, + ServerStreams: true, + }, }, Metadata: "cortex-axon-agent.proto", } diff --git a/sdks/go/axon_agent.go b/sdks/go/axon_agent.go index 87e7d967..bd930f0f 100644 --- a/sdks/go/axon_agent.go +++ b/sdks/go/axon_agent.go @@ -36,6 +36,9 @@ type Agent struct { logger *zap.Logger sleepOnError time.Duration done chan struct{} + + preDispatch PreDispatchHook + postDispatch PostDispatchHook } // NewAxonAgent creates a new AxonAgent with the specified options. You @@ -58,6 +61,8 @@ func NewAxonAgent(options ...Option) *Agent { logger: logger, sleepOnError: ao.sleepOnError, done: make(chan struct{}), + preDispatch: ao.preDispatch, + postDispatch: ao.postDispatch, } a.logger = logger @@ -247,6 +252,23 @@ func (a *Agent) Run(ctx context.Context) error { runningHandlers := sync.WaitGroup{} + // Failure events stop after the last handler finishes, so a hook that + // retries from a handler still has a live agent connection. + eventsCtx, stopEvents := context.WithCancel(context.Background()) + eventsDone := make(chan struct{}) + if a.postDispatch != nil { + go func() { + defer close(eventsDone) + a.runFailureEvents(eventsCtx) + }() + } else { + close(eventsDone) + } + defer func() { + stopEvents() + <-eventsDone + }() + go func() { select { case <-ctx.Done(): @@ -287,6 +309,7 @@ func (a *Agent) Run(ctx context.Context) error { err = stream.Send(&pb.DispatchRequest{ DispatchId: a.DispatchId, ClientVersion: version.Client, + FailureEvents: a.postDispatch != nil, }) if err != nil { a.logger.Error("failed to send registration id", zap.Error(err)) @@ -399,46 +422,77 @@ func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInv return } - done := make(chan struct{}) + info := newInvocationInfo(invoke) + a.runPreDispatch(info) report := &pb.ReportInvocationRequest{ HandlerInvoke: invoke, StartClientTimestamp: timestamppb.Now(), } - - go func() { - apiStub := a.client.api() - - wrapped := zapcore.RegisterHooks(a.logger.Core(), func(entry zapcore.Entry) error { - report.Logs = append(report.Logs, &pb.Log{ - Level: entry.Level.CapitalString(), - Message: entry.Message, - Timestamp: timestamppb.New(entry.Time), - }) - return nil + start := time.Now() + + // The handler keeps running after a timeout, so it must not touch the + // report: logs go through a lock and the outcome over a channel. + var logsMu sync.Mutex + var logs []*pb.Log + wrapped := zapcore.RegisterHooks(a.logger.Core(), func(entry zapcore.Entry) error { + logsMu.Lock() + defer logsMu.Unlock() + logs = append(logs, &pb.Log{ + Level: entry.Level.CapitalString(), + Message: entry.Message, + Timestamp: timestamppb.New(entry.Time), }) + return nil + }) - loggerFromCore := zap.New(wrapped) - - handlerContext := NewHandlerContext(invoke, ctx, apiStub, loggerFromCore) - result, duration, err := a.executeHandlerWithRecover(handlerInfo, handlerContext) - report.DurationMs = int32(duration.Milliseconds()) + loggerFromCore := zap.New(wrapped) + hctx := NewHandlerContext(invoke, ctx, a.client.api(), loggerFromCore) - if err != nil { - a.setReportError(report, "unexpected", err) - } else { - a.setReportResult(report, result) - } - - close(done) + type handlerOutcome struct { + result any + duration time.Duration + err error + } + done := make(chan handlerOutcome, 1) + go func() { + result, duration, err := a.executeHandlerWithRecover(handlerInfo, hctx) + done <- handlerOutcome{result, duration, err} }() + var failure *Failure select { - case <-done: + case out := <-done: + report.DurationMs = int32(out.duration.Milliseconds()) + if out.err != nil { + a.setReportError(report, "unexpected", out.err) + failure = &Failure{ + Reason: pb.FailureReason_HANDLER_ERROR, + Source: pb.FailureSource_APPLICATION, + Message: out.err.Error(), + } + } else { + a.setReportResult(report, out.result) + } case <-ctx.Done(): + report.DurationMs = int32(time.Since(start).Milliseconds()) a.setReportError(report, "timeout", nil) + failure = &Failure{ + Reason: pb.FailureReason_TIMEOUT, + Source: pb.FailureSource_APPLICATION, + Message: fmt.Sprintf("handler did not finish within %dms", invoke.TimeoutMs), + } + if hc, ok := hctx.(*handlerContext); ok && hc.calls.blockedOnApi() { + failure.Source = pb.FailureSource_RUNTIME + failure.Message += " while waiting on a Cortex API call" + } } + logsMu.Lock() + report.Logs = logs + logs = nil + logsMu.Unlock() + if report.GetError() == nil { a.logger.Debug( "handler executed", @@ -457,6 +511,10 @@ func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInv if err != nil { a.logger.Error("failed to report invocation", zap.Error(err)) } + + // After the report, so a slow hook can't make the agent decide this + // invocation was never answered and report the failure a second time. + a.runPostDispatch(info, failure) } func (a *Agent) executeHandlerWithRecover(handler *handlerInfo, ctx HandlerContext) (result any, d time.Duration, err error) { diff --git a/sdks/go/axon_agent_test.go b/sdks/go/axon_agent_test.go index a6b965d6..adfe0638 100644 --- a/sdks/go/axon_agent_test.go +++ b/sdks/go/axon_agent_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "io" + "sync/atomic" "testing" "time" @@ -103,10 +104,11 @@ func TestInvokeHandlerWithTimeoutTriggered(t *testing.T) { controller := gomock.NewController(t) defer controller.Finish() - called := false + // The handler outlives the timeout, so it runs concurrently with the check. + var called atomic.Bool theHandler := func(ctx HandlerContext) error { - called = true + called.Store(true) time.Sleep(time.Millisecond * 10) return nil } @@ -117,7 +119,7 @@ func TestInvokeHandlerWithTimeoutTriggered(t *testing.T) { require.Equal(t, "timeout", reportedErr.Code) }) - require.True(t, called) + require.True(t, called.Load()) } @@ -237,6 +239,16 @@ func createAgent(controller *gomock.Controller) (*Agent, *mockGrpcClient) { type agentCallback func(*mockGrpcClient) func executeHandlerHelper(t *testing.T, controller *gomock.Controller, handler any, timeoutMs int, reportCallback func(*pb.ReportInvocationRequest), callback ...agentCallback) error { + return executeHandlerWithAgent(t, controller, handler, timeoutMs, reportCallback, func(_ *Agent, mock *mockGrpcClient) { + if len(callback) > 0 { + callback[0](mock) + } + }) +} + +// executeHandlerWithAgent runs one invocation of handler. configure can set +// up the agent and mocks before it runs. +func executeHandlerWithAgent(t *testing.T, controller *gomock.Controller, handler any, timeoutMs int, reportCallback func(*pb.ReportInvocationRequest), configure func(*Agent, *mockGrpcClient)) error { agent, mock := createAgent(controller) id := fmt.Sprintf("%d", time.Now().UnixNano()) @@ -257,6 +269,7 @@ func executeHandlerHelper(t *testing.T, controller *gomock.Controller, handler a }, ) + mock.lastStream = stream mock.agentStub.EXPECT().RegisterHandler(gomock.Any(), gomock.Any()).Return(&pb.RegisterHandlerResponse{Id: id}, nil) mock.agentStub.EXPECT().Dispatch(gomock.Any(), gomock.Any()).AnyTimes().Return(stream, nil) mock.agentStub.EXPECT().ReportInvocation(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( @@ -267,9 +280,7 @@ func executeHandlerHelper(t *testing.T, controller *gomock.Controller, handler a return &pb.ReportInvocationResponse{}, nil }) - if len(callback) > 0 { - callback[0](mock) - } + configure(agent, mock) options := []RegisterHandlerOption{ WithInvokeOption(pb.HandlerInvokeType_RUN_NOW, ""), @@ -302,6 +313,8 @@ func executeHandlerHelper(t *testing.T, controller *gomock.Controller, handler a type mockGrpcClient struct { apiStub *mock_axon.MockCortexApiClient agentStub *mock_axon.MockAxonAgentClient + + lastStream *mockBidiClient } func (m *mockGrpcClient) api() pb.CortexApiClient { @@ -318,6 +331,7 @@ type mockBidiClient struct { messages []*pb.DispatchMessage messageIndex int eofOnFinish bool + sent []*pb.DispatchRequest } func newBidiClient(messages ...*pb.DispatchMessage) *mockBidiClient { @@ -346,7 +360,8 @@ func (m *mockBidiClient) Recv() (*pb.DispatchMessage, error) { return msg, nil } -func (m *mockBidiClient) Send(*pb.DispatchRequest) error { +func (m *mockBidiClient) Send(req *pb.DispatchRequest) error { + m.sent = append(m.sent, req) if m.IsDone() && m.eofOnFinish { return io.EOF } diff --git a/sdks/go/handler_context.go b/sdks/go/handler_context.go index 4e4fd632..cdbe9280 100644 --- a/sdks/go/handler_context.go +++ b/sdks/go/handler_context.go @@ -2,9 +2,11 @@ package axon import ( "context" + "sync/atomic" pb "github.com/cortexapps/axon-go/.generated/proto/github.com/cortexapps/axon" "go.uber.org/zap" + "google.golang.org/grpc" ) type handlerContextKey string @@ -22,7 +24,36 @@ type HandlerContext interface { type handlerContext struct { context.Context - args map[string]string + args map[string]string + calls *apiCallTracker +} + +// apiCallTracker records whether a Cortex API call was running when the +// handler's deadline hit, which makes the timeout Axon's fault, not the +// handler's. +type apiCallTracker struct { + inFlight atomic.Int32 + cutOff atomic.Bool +} + +func (t *apiCallTracker) blockedOnApi() bool { + return t.inFlight.Load() > 0 || t.cutOff.Load() +} + +type trackedApiClient struct { + pb.CortexApiClient + calls *apiCallTracker +} + +func (c trackedApiClient) Call(ctx context.Context, in *pb.CallRequest, opts ...grpc.CallOption) (*pb.CallResponse, error) { + c.calls.inFlight.Add(1) + defer c.calls.inFlight.Add(-1) + resp, err := c.CortexApiClient.Call(ctx, in, opts...) + if err != nil && ctx.Err() != nil { + // Set before inFlight drops, so a timeout check never sees neither. + c.calls.cutOff.Store(true) + } + return resp, err } func (h *handlerContext) Args() map[string]string { @@ -54,11 +85,13 @@ func NewHandlerContext(invoke *pb.DispatchHandlerInvoke, ctx context.Context, ap name := invoke.HandlerName logger = logger.With(zap.String("handler-name", name)) + calls := &apiCallTracker{} ctx = context.WithValue(ctx, logKey, logger) - ctx = context.WithValue(ctx, apiKey, api) + ctx = context.WithValue(ctx, apiKey, pb.CortexApiClient(trackedApiClient{CortexApiClient: api, calls: calls})) return &handlerContext{ Context: ctx, args: invoke.Args, + calls: calls, } } diff --git a/sdks/go/hooks.go b/sdks/go/hooks.go new file mode 100644 index 00000000..c3e8234c --- /dev/null +++ b/sdks/go/hooks.go @@ -0,0 +1,203 @@ +package axon + +import ( + "context" + "errors" + "fmt" + "time" + + pb "github.com/cortexapps/axon-go/.generated/proto/github.com/cortexapps/axon" + "github.com/cortexapps/axon-go/version" + "go.uber.org/zap" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// InvocationInfo identifies one invocation for the dispatch hooks. For a +// webhook, Args holds the inbound message: "body", "content-type" and "url". +type InvocationInfo struct { + InvocationID string + HandlerName string + InvokeType pb.HandlerInvokeType + Args map[string]string + // Attempt is 1 for the first delivery and counts up with each retry. + Attempt int + + invoke *pb.DispatchHandlerInvoke +} + +func newInvocationInfo(invoke *pb.DispatchHandlerInvoke) InvocationInfo { + return InvocationInfo{ + InvocationID: invoke.InvocationId, + HandlerName: invoke.HandlerName, + InvokeType: invoke.Reason, + Args: invoke.Args, + Attempt: max(int(invoke.Attempt), 1), + invoke: invoke, + } +} + +// Failure describes why an invocation failed. Source says whether handler +// code failed (APPLICATION) or Axon did (RUNTIME). +type Failure struct { + Reason pb.FailureReason + Source pb.FailureSource + Message string + // CallerRejected means whoever sent the invocation was told it was not + // accepted (a webhook sender got a 503), so it may arrive again. + CallerRejected bool +} + +func failureFromProto(f *pb.InvocationFailure) *Failure { + return &Failure{ + Reason: f.GetReason(), + Source: f.GetSource(), + Message: f.GetMessage(), + CallerRejected: f.GetCallerRejected(), + } +} + +// PostAction is what a PostDispatchHook wants done with a failed invocation. +type PostAction int + +const ( + // PostActionIgnore leaves the failure as it is. + PostActionIgnore PostAction = iota + // PostActionRetry queues the invocation again as its next attempt. + PostActionRetry +) + +// PreDispatchHook runs before each invocation's handler. +type PreDispatchHook func(ctx context.Context, info InvocationInfo) + +// PostDispatchHook runs after each invocation, with a nil failure when it +// succeeded. It also runs for failures the agent detects on its own, such as +// an invocation that never reached this process, so it sees every failed +// message. Its return value only matters for a failure. +type PostDispatchHook func(ctx context.Context, info InvocationInfo, failure *Failure) PostAction + +// WithPreDispatchHook sets a hook that runs before every handler. +func WithPreDispatchHook(hook PreDispatchHook) Option { + return func(a *agentOptions) { + a.preDispatch = hook + } +} + +// WithPostDispatchHook sets a hook that runs after every invocation and for +// every failure, including ones the agent detects. +func WithPostDispatchHook(hook PostDispatchHook) Option { + return func(a *agentOptions) { + a.postDispatch = hook + } +} + +func (a *Agent) runPreDispatch(info InvocationInfo) { + if a.preDispatch == nil { + return + } + defer func() { + if r := recover(); r != nil { + a.logger.Error("panic in pre-dispatch hook", zap.Any("panic", r)) + } + }() + a.preDispatch(context.Background(), info) +} + +func (a *Agent) runPostDispatch(info InvocationInfo, failure *Failure) { + if a.postDispatch == nil { + return + } + if a.callPostDispatch(info, failure) == PostActionRetry && failure != nil { + a.retryInvocation(info) + } +} + +func (a *Agent) callPostDispatch(info InvocationInfo, failure *Failure) (action PostAction) { + defer func() { + if r := recover(); r != nil { + a.logger.Error("panic in post-dispatch hook", zap.Any("panic", r)) + action = PostActionIgnore + } + }() + return a.postDispatch(context.Background(), info, failure) +} + +func (a *Agent) retryInvocation(info InvocationInfo) { + stub := a.client.agent() + if stub == nil { + a.logger.Error("cannot retry invocation, no agent connection", + zap.String("invocation-id", info.InvocationID)) + return + } + resp, err := stub.RetryInvocation(context.Background(), &pb.RetryInvocationRequest{Invoke: info.invoke}) + if status.Code(err) == codes.Unimplemented { + a.logger.Warn("agent does not support retries, upgrade the agent", + zap.String("invocation-id", info.InvocationID)) + return + } + if err != nil { + a.logger.Error("failed to retry invocation", + zap.String("invocation-id", info.InvocationID), zap.Error(err)) + return + } + if resp.GetError() != nil { + a.logger.Warn("agent rejected retry", + zap.String("invocation-id", info.InvocationID), + zap.String("error", resp.GetError().GetMessage())) + return + } + a.logger.Info("retrying invocation", + zap.String("invocation-id", info.InvocationID), + zap.String("retry-invocation-id", resp.GetInvocationId()), + zap.Int("attempt", info.Attempt+1)) +} + +var errEventsUnsupported = errors.New("agent does not support failure events") + +// runFailureEvents reads the agent's failure events until ctx is done. It has +// its own stream so it keeps reading while every handler is busy. +func (a *Agent) runFailureEvents(ctx context.Context) { + for ctx.Err() == nil { + err := a.readFailureEvents(ctx) + if errors.Is(err, errEventsUnsupported) { + a.logger.Warn("agent does not send failure events, upgrade the agent to receive them") + return + } + if ctx.Err() != nil { + return + } + a.logger.Warn("failure event stream ended, reconnecting in "+sleepErrorWait.String(), zap.Error(err)) + select { + case <-ctx.Done(): + case <-time.After(sleepErrorWait): + } + } +} + +func (a *Agent) readFailureEvents(ctx context.Context) error { + stub := a.client.agent() + if stub == nil { + return fmt.Errorf("failed to create agent connection") + } + stream, err := stub.DispatchEvents(ctx, &pb.DispatchRequest{ + DispatchId: a.DispatchId, + ClientVersion: version.Client, + FailureEvents: true, + }) + if err != nil { + return err + } + for { + ev, err := stream.Recv() + if status.Code(err) == codes.Unimplemented { + return errEventsUnsupported + } + if err != nil { + return err + } + if ev.GetInvoke() == nil { + continue + } + a.runPostDispatch(newInvocationInfo(ev.GetInvoke()), failureFromProto(ev.GetFailure())) + } +} diff --git a/sdks/go/hooks_test.go b/sdks/go/hooks_test.go new file mode 100644 index 00000000..bddcbf7b --- /dev/null +++ b/sdks/go/hooks_test.go @@ -0,0 +1,284 @@ +package axon + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + pb "github.com/cortexapps/axon-go/.generated/proto/github.com/cortexapps/axon" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type hookCall struct { + info InvocationInfo + failure *Failure +} + +// hookRecorder is a post-dispatch hook that records each call. +type hookRecorder struct { + mu sync.Mutex + calls []hookCall + action PostAction + seen chan struct{} +} + +func newHookRecorder(action PostAction) *hookRecorder { + return &hookRecorder{action: action, seen: make(chan struct{}, 10)} +} + +func (r *hookRecorder) hook(_ context.Context, info InvocationInfo, failure *Failure) PostAction { + r.mu.Lock() + r.calls = append(r.calls, hookCall{info, failure}) + r.mu.Unlock() + r.seen <- struct{}{} + return r.action +} + +func (r *hookRecorder) only(t *testing.T) hookCall { + r.mu.Lock() + defer r.mu.Unlock() + require.Len(t, r.calls, 1) + return r.calls[0] +} + +// eventStream is a DispatchEvents stream that delivers events and then +// blocks until the agent stops reading. +type eventStream struct { + grpc.ClientStream + ctx context.Context + events chan *pb.DispatchEvent + err error +} + +func (s *eventStream) Recv() (*pb.DispatchEvent, error) { + if s.err != nil { + return nil, s.err + } + select { + case ev := <-s.events: + return ev, nil + case <-s.ctx.Done(): + return nil, s.ctx.Err() + } +} + +func expectEventStream(mock *mockGrpcClient, events ...*pb.DispatchEvent) { + mock.agentStub.EXPECT().DispatchEvents(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn( + func(ctx context.Context, req *pb.DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[pb.DispatchEvent], error) { + ch := make(chan *pb.DispatchEvent, len(events)) + for _, ev := range events { + ch <- ev + } + events = nil + return &eventStream{ctx: ctx, events: ch}, nil + }) +} + +func runWithPostHook(t *testing.T, handler any, timeoutMs int, rec *hookRecorder, setup ...func(*mockGrpcClient)) { + controller := gomock.NewController(t) + defer controller.Finish() + executeHandlerWithAgent(t, controller, handler, timeoutMs, nil, func(agent *Agent, mock *mockGrpcClient) { + agent.postDispatch = rec.hook + expectEventStream(mock) + for _, s := range setup { + s(mock) + } + }) +} + +func TestPostDispatchHookOnSuccess(t *testing.T) { + rec := newHookRecorder(PostActionRetry) + runWithPostHook(t, func(ctx HandlerContext) error { return nil }, 0, rec) + + call := rec.only(t) + require.Nil(t, call.failure) + require.Equal(t, "func1", call.info.HandlerName) + require.Equal(t, pb.HandlerInvokeType_RUN_NOW, call.info.InvokeType) + require.Equal(t, 1, call.info.Attempt, "an agent that doesn't send attempt means the first") + // Retry on success is ignored: no RetryInvocation expectation is set. +} + +func TestPostDispatchHookOnHandlerError(t *testing.T) { + rec := newHookRecorder(PostActionIgnore) + runWithPostHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec) + + failure := rec.only(t).failure + require.NotNil(t, failure) + require.Equal(t, pb.FailureReason_HANDLER_ERROR, failure.Reason) + require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) + require.Equal(t, "ohnoes", failure.Message) +} + +func TestPostDispatchHookOnPanic(t *testing.T) { + rec := newHookRecorder(PostActionIgnore) + runWithPostHook(t, func(ctx HandlerContext) error { panic("ohnoes") }, 0, rec) + + failure := rec.only(t).failure + require.Equal(t, pb.FailureReason_HANDLER_ERROR, failure.Reason) + require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) +} + +func TestPostDispatchHookOnHandlerTimeout(t *testing.T) { + rec := newHookRecorder(PostActionIgnore) + runWithPostHook(t, func(ctx HandlerContext) error { + time.Sleep(50 * time.Millisecond) + return nil + }, 1, rec) + + failure := rec.only(t).failure + require.Equal(t, pb.FailureReason_TIMEOUT, failure.Reason) + require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) +} + +func TestPostDispatchHookTimeoutDuringApiCallIsRuntime(t *testing.T) { + rec := newHookRecorder(PostActionIgnore) + runWithPostHook(t, func(ctx HandlerContext) error { + _, err := ctx.CortexJsonApiCall("GET", "/api/v1/slow", "") + return err + }, 20, rec, func(mock *mockGrpcClient) { + mock.apiStub.EXPECT().Call(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, req *pb.CallRequest, opts ...grpc.CallOption) (*pb.CallResponse, error) { + <-ctx.Done() + return nil, status.FromContextError(ctx.Err()).Err() + }) + }) + + failure := rec.only(t).failure + require.Equal(t, pb.FailureReason_TIMEOUT, failure.Reason) + require.Equal(t, pb.FailureSource_RUNTIME, failure.Source) +} + +func TestPostDispatchHookRetry(t *testing.T) { + rec := newHookRecorder(PostActionRetry) + retried := make(chan *pb.RetryInvocationRequest, 1) + runWithPostHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec, func(mock *mockGrpcClient) { + mock.agentStub.EXPECT().RetryInvocation(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, req *pb.RetryInvocationRequest, opts ...grpc.CallOption) (*pb.RetryInvocationResponse, error) { + retried <- req + return &pb.RetryInvocationResponse{InvocationId: "retry-1"}, nil + }) + }) + + req := <-retried + require.Equal(t, "func1", req.Invoke.HandlerName) +} + +func TestPostDispatchHookPanicIsContained(t *testing.T) { + controller := gomock.NewController(t) + defer controller.Finish() + + reported := false + executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, + func(*pb.ReportInvocationRequest) { reported = true }, + func(agent *Agent, mock *mockGrpcClient) { + agent.postDispatch = func(context.Context, InvocationInfo, *Failure) PostAction { panic("hook bug") } + expectEventStream(mock) + }) + require.True(t, reported) +} + +func TestPreDispatchHookRunsBeforeHandler(t *testing.T) { + controller := gomock.NewController(t) + defer controller.Finish() + + var order []string + executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { + order = append(order, "handler") + return nil + }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { + agent.preDispatch = func(_ context.Context, info InvocationInfo) { + require.Equal(t, "func1", info.HandlerName) + order = append(order, "pre") + } + }) + require.Equal(t, []string{"pre", "handler"}, order) +} + +func TestFailureEventReachesPostDispatchHook(t *testing.T) { + rec := newHookRecorder(PostActionRetry) + retried := make(chan *pb.RetryInvocationRequest, 1) + + ev := &pb.DispatchEvent{ + Invoke: &pb.DispatchHandlerInvoke{ + InvocationId: "lost-1", + HandlerName: "func1", + Reason: pb.HandlerInvokeType_WEBHOOK, + Args: map[string]string{"body": "payload"}, + Attempt: 2, + }, + Failure: &pb.InvocationFailure{ + Reason: pb.FailureReason_NO_RESOURCES, + Source: pb.FailureSource_RUNTIME, + Message: "dispatch queue is full", + CallerRejected: true, + }, + } + + controller := gomock.NewController(t) + defer controller.Finish() + executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { + // Hold the run open until the event has been handled. + <-rec.seen + return nil + }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { + agent.postDispatch = rec.hook + expectEventStream(mock, ev) + mock.agentStub.EXPECT().RetryInvocation(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, req *pb.RetryInvocationRequest, opts ...grpc.CallOption) (*pb.RetryInvocationResponse, error) { + retried <- req + return &pb.RetryInvocationResponse{InvocationId: "retry-1"}, nil + }) + }) + + rec.mu.Lock() + first := rec.calls[0] + rec.mu.Unlock() + require.Equal(t, "lost-1", first.info.InvocationID) + require.Equal(t, "payload", first.info.Args["body"]) + require.Equal(t, 2, first.info.Attempt) + require.Equal(t, pb.FailureReason_NO_RESOURCES, first.failure.Reason) + require.True(t, first.failure.CallerRejected) + require.Equal(t, "lost-1", (<-retried).Invoke.InvocationId) +} + +func TestFailureEventsStopOnOlderAgent(t *testing.T) { + controller := gomock.NewController(t) + defer controller.Finish() + + executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { + time.Sleep(50 * time.Millisecond) + return nil + }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { + agent.postDispatch = func(context.Context, InvocationInfo, *Failure) PostAction { return PostActionIgnore } + // Times(1): an agent without the RPC must not be redialed in a loop. + mock.agentStub.EXPECT().DispatchEvents(gomock.Any(), gomock.Any()).Times(1).Return( + &eventStream{err: status.Error(codes.Unimplemented, "unknown method DispatchEvents")}, nil) + }) +} + +func TestDispatchRequestAdvertisesFailureEvents(t *testing.T) { + for _, withHook := range []bool{false, true} { + t.Run(fmt.Sprint(withHook), func(t *testing.T) { + controller := gomock.NewController(t) + defer controller.Finish() + + var sent *mockBidiClient + executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { return nil }, 0, nil, + func(agent *Agent, mock *mockGrpcClient) { + if withHook { + agent.postDispatch = func(context.Context, InvocationInfo, *Failure) PostAction { return PostActionIgnore } + expectEventStream(mock) + } + sent = mock.lastStream + }) + require.Len(t, sent.sent, 1) + require.Equal(t, withHook, sent.sent[0].FailureEvents) + }) + } +} diff --git a/sdks/go/mock_axon/mocks.go b/sdks/go/mock_axon/mocks.go index 5672198c..40fd1196 100644 --- a/sdks/go/mock_axon/mocks.go +++ b/sdks/go/mock_axon/mocks.go @@ -106,6 +106,26 @@ func (mr *MockAxonAgentClientMockRecorder) Dispatch(ctx any, opts ...any) *gomoc return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Dispatch", reflect.TypeOf((*MockAxonAgentClient)(nil).Dispatch), varargs...) } +// DispatchEvents mocks base method. +func (m *MockAxonAgentClient) DispatchEvents(ctx context.Context, in *axon.DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[axon.DispatchEvent], error) { + m.ctrl.T.Helper() + varargs := []any{ctx, in} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "DispatchEvents", varargs...) + ret0, _ := ret[0].(grpc.ServerStreamingClient[axon.DispatchEvent]) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DispatchEvents indicates an expected call of DispatchEvents. +func (mr *MockAxonAgentClientMockRecorder) DispatchEvents(ctx, in any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, in}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchEvents", reflect.TypeOf((*MockAxonAgentClient)(nil).DispatchEvents), varargs...) +} + // GetHandlerHistory mocks base method. func (m *MockAxonAgentClient) GetHandlerHistory(ctx context.Context, in *axon.GetHandlerHistoryRequest, opts ...grpc.CallOption) (*axon.GetHandlerHistoryResponse, error) { m.ctrl.T.Helper() @@ -186,6 +206,26 @@ func (mr *MockAxonAgentClientMockRecorder) ReportInvocation(ctx, in any, opts .. return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReportInvocation", reflect.TypeOf((*MockAxonAgentClient)(nil).ReportInvocation), varargs...) } +// RetryInvocation mocks base method. +func (m *MockAxonAgentClient) RetryInvocation(ctx context.Context, in *axon.RetryInvocationRequest, opts ...grpc.CallOption) (*axon.RetryInvocationResponse, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, in} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "RetryInvocation", varargs...) + ret0, _ := ret[0].(*axon.RetryInvocationResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RetryInvocation indicates an expected call of RetryInvocation. +func (mr *MockAxonAgentClientMockRecorder) RetryInvocation(ctx, in any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, in}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RetryInvocation", reflect.TypeOf((*MockAxonAgentClient)(nil).RetryInvocation), varargs...) +} + // UnregisterHandler mocks base method. func (m *MockAxonAgentClient) UnregisterHandler(ctx context.Context, in *axon.UnregisterHandlerRequest, opts ...grpc.CallOption) (*axon.UnregisterHandlerResponse, error) { m.ctrl.T.Helper() diff --git a/sdks/go/options.go b/sdks/go/options.go index bb862114..e55bb148 100644 --- a/sdks/go/options.go +++ b/sdks/go/options.go @@ -28,6 +28,8 @@ type agentOptions struct { sleepOnError time.Duration version string maxReceiveMessageSize int + preDispatch PreDispatchHook + postDispatch PostDispatchHook } func resolveMaxReceiveMessageSize() int { From de9309c2f33f05f4fc0943142ed59ae492fa0378 Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Wed, 11 Nov 2026 10:47:45 +1100 Subject: [PATCH 2/5] Failure hooks for the Python SDK; don't run invocations the agent gave up on Python gets pre_dispatch and post_dispatch on AxonClient, matching Go. Failure events are read on their own thread and run on their own pool, so they arrive even while every handler slot is busy. At a handler's timeout the SDK now reports the timeout itself and drops whatever the handler does afterwards, since a Python thread can't be stopped. The agent now sends each invocation an absolute deadline: when it stops waiting and reports the failure itself. A saturated client used to read invocations out of the stream buffer long after the agent had reported them as failed, and run them anyway, so a failure hook replaying them would have processed them twice. Both SDKs now skip an invocation past its deadline and cap a handler's timeout at it, classing that timeout as RUNTIME. Co-Authored-By: Claude Opus 5.5 --- agent/proto/cortex-axon-agent.proto | 4 + agent/server/server.go | 6 +- .../cortexapps/axon/cortex-axon-agent.pb.go | 105 ++--- sdks/go/README.md | 30 ++ sdks/go/axon_agent.go | 49 ++- sdks/go/hooks_test.go | 28 ++ sdks/python/README.md | 63 +++ sdks/python/cortex_axon/__init__.py | 5 + sdks/python/cortex_axon/axon_client.py | 266 ++++++++++++- sdks/python/cortex_axon/hooks.py | 85 ++++ sdks/python/example/main.py | 16 + sdks/python/tests/test_hooks.py | 370 ++++++++++++++++++ 12 files changed, 964 insertions(+), 63 deletions(-) create mode 100644 sdks/python/cortex_axon/hooks.py create mode 100644 sdks/python/tests/test_hooks.py diff --git a/agent/proto/cortex-axon-agent.proto b/agent/proto/cortex-axon-agent.proto index 09dd16b8..c4ca4266 100644 --- a/agent/proto/cortex-axon-agent.proto +++ b/agent/proto/cortex-axon-agent.proto @@ -125,6 +125,10 @@ message DispatchHandlerInvoke{ map args = 20; // 1 for the first attempt; 0 means the agent didn't say, so treat it as 1. int32 attempt = 30; + // When the agent stops waiting for this invocation's result and reports + // the failure itself. A client must not start a handler after it, since + // the failure has already been reported. Unset from older agents. + google.protobuf.Timestamp deadline = 31; } // diff --git a/agent/server/server.go b/agent/server/server.go index 7f80f90b..d8cd6627 100644 --- a/agent/server/server.go +++ b/agent/server/server.go @@ -205,6 +205,10 @@ func (s *AxonAgent) sendInvocations(dispatchId string, stream pb.AxonAgent_Dispa now := time.Now() msg := invoke.ToDispatchInvoke() + // The client may not read this for a while if every handler is + // busy, so it gets an absolute deadline rather than only a timeout. + deadline := now.Add(time.Duration(msg.TimeoutMs)*time.Millisecond + reportGrace) + msg.Deadline = timestamppb.New(deadline) dispatchMessage := &pb.DispatchMessage{ Type: pb.DispatchMessageType_DISPATCH_MESSAGE_INVOKE, @@ -224,7 +228,7 @@ func (s *AxonAgent) sendInvocations(dispatchId string, stream pb.AxonAgent_Dispa // The grace lets a client that enforces the timeout itself // report it first, so the failure is described once, by the // side that saw it. - <-time.After(time.Duration(msg.TimeoutMs)*time.Millisecond + reportGrace) + <-time.After(time.Until(deadline)) message := fmt.Sprintf("handler did not report a result within %dms", msg.TimeoutMs) completed := s.completeInvocation(context.Background(), &pb.ReportInvocationRequest{ HandlerInvoke: msg, diff --git a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go index 5888e0df..36684fc1 100644 --- a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go +++ b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go @@ -877,7 +877,11 @@ type DispatchHandlerInvoke struct { Reason HandlerInvokeType `protobuf:"varint,11,opt,name=reason,proto3,enum=cortex.axon.HandlerInvokeType" json:"reason,omitempty"` Args map[string]string `protobuf:"bytes,20,rep,name=args,proto3" json:"args,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` // 1 for the first attempt; 0 means the agent didn't say, so treat it as 1. - Attempt int32 `protobuf:"varint,30,opt,name=attempt,proto3" json:"attempt,omitempty"` + Attempt int32 `protobuf:"varint,30,opt,name=attempt,proto3" json:"attempt,omitempty"` + // When the agent stops waiting for this invocation's result and reports + // the failure itself. A client must not start a handler after it, since + // the failure has already been reported. Unset from older agents. + Deadline *timestamppb.Timestamp `protobuf:"bytes,31,opt,name=deadline,proto3" json:"deadline,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -968,6 +972,13 @@ func (x *DispatchHandlerInvoke) GetAttempt() int32 { return 0 } +func (x *DispatchHandlerInvoke) GetDeadline() *timestamppb.Timestamp { + if x != nil { + return x.Deadline + } + return nil +} + type InvocationFailure struct { state protoimpl.MessageState `protogen:"open.v1"` Reason FailureReason `protobuf:"varint,1,opt,name=reason,proto3,enum=cortex.axon.FailureReason" json:"reason,omitempty"` @@ -1738,7 +1749,7 @@ const file_cortex_axon_agent_proto_rawDesc = "" + "\x04type\x18\x01 \x01(\x0e2 .cortex.axon.DispatchMessageTypeR\x04type\x12<\n" + "\x06invoke\x18\n" + " \x01(\v2\".cortex.axon.DispatchHandlerInvokeH\x00R\x06invokeB\t\n" + - "\amessage\"\x8b\x03\n" + + "\amessage\"\xc3\x03\n" + "\x15DispatchHandlerInvoke\x12#\n" + "\rinvocation_id\x18\x01 \x01(\tR\finvocationId\x12\x1f\n" + "\vdispatch_id\x18\x02 \x01(\tR\n" + @@ -1751,7 +1762,8 @@ const file_cortex_axon_agent_proto_rawDesc = "" + " \x01(\x05R\ttimeoutMs\x126\n" + "\x06reason\x18\v \x01(\x0e2\x1e.cortex.axon.HandlerInvokeTypeR\x06reason\x12@\n" + "\x04args\x18\x14 \x03(\v2,.cortex.axon.DispatchHandlerInvoke.ArgsEntryR\x04args\x12\x18\n" + - "\aattempt\x18\x1e \x01(\x05R\aattempt\x1a7\n" + + "\aattempt\x18\x1e \x01(\x05R\aattempt\x126\n" + + "\bdeadline\x18\x1f \x01(\v2\x1a.google.protobuf.TimestampR\bdeadline\x1a7\n" + "\tArgsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xbe\x01\n" + @@ -1901,49 +1913,50 @@ var file_cortex_axon_agent_proto_depIdxs = []int32{ 15, // 10: cortex.axon.DispatchMessage.invoke:type_name -> cortex.axon.DispatchHandlerInvoke 0, // 11: cortex.axon.DispatchHandlerInvoke.reason:type_name -> cortex.axon.HandlerInvokeType 27, // 12: cortex.axon.DispatchHandlerInvoke.args:type_name -> cortex.axon.DispatchHandlerInvoke.ArgsEntry - 2, // 13: cortex.axon.InvocationFailure.reason:type_name -> cortex.axon.FailureReason - 3, // 14: cortex.axon.InvocationFailure.source:type_name -> cortex.axon.FailureSource - 15, // 15: cortex.axon.DispatchEvent.invoke:type_name -> cortex.axon.DispatchHandlerInvoke - 16, // 16: cortex.axon.DispatchEvent.failure:type_name -> cortex.axon.InvocationFailure - 15, // 17: cortex.axon.RetryInvocationRequest.invoke:type_name -> cortex.axon.DispatchHandlerInvoke - 28, // 18: cortex.axon.RetryInvocationResponse.error:type_name -> cortex.axon.Error - 29, // 19: cortex.axon.Log.timestamp:type_name -> google.protobuf.Timestamp - 15, // 20: cortex.axon.ReportInvocationRequest.handler_invoke:type_name -> cortex.axon.DispatchHandlerInvoke - 29, // 21: cortex.axon.ReportInvocationRequest.start_client_timestamp:type_name -> google.protobuf.Timestamp - 23, // 22: cortex.axon.ReportInvocationRequest.result:type_name -> cortex.axon.InvokeResult - 28, // 23: cortex.axon.ReportInvocationRequest.error:type_name -> cortex.axon.Error - 20, // 24: cortex.axon.ReportInvocationRequest.logs:type_name -> cortex.axon.Log - 28, // 25: cortex.axon.ReportInvocationResponse.error:type_name -> cortex.axon.Error - 29, // 26: cortex.axon.GetHandlerHistoryRequest.start_time:type_name -> google.protobuf.Timestamp - 29, // 27: cortex.axon.GetHandlerHistoryRequest.end_time:type_name -> google.protobuf.Timestamp - 29, // 28: cortex.axon.HandlerExecution.publish_server_timestamp:type_name -> google.protobuf.Timestamp - 29, // 29: cortex.axon.HandlerExecution.receive_server_timestamp:type_name -> google.protobuf.Timestamp - 29, // 30: cortex.axon.HandlerExecution.start_client_timestamp:type_name -> google.protobuf.Timestamp - 28, // 31: cortex.axon.HandlerExecution.error:type_name -> cortex.axon.Error - 20, // 32: cortex.axon.HandlerExecution.logs:type_name -> cortex.axon.Log - 28, // 33: cortex.axon.GetHandlerHistoryResponse.error:type_name -> cortex.axon.Error - 25, // 34: cortex.axon.GetHandlerHistoryResponse.history:type_name -> cortex.axon.HandlerExecution - 4, // 35: cortex.axon.AxonAgent.RegisterHandler:input_type -> cortex.axon.RegisterHandlerRequest - 8, // 36: cortex.axon.AxonAgent.UnregisterHandler:input_type -> cortex.axon.UnregisterHandlerRequest - 10, // 37: cortex.axon.AxonAgent.ListHandlers:input_type -> cortex.axon.ListHandlersRequest - 24, // 38: cortex.axon.AxonAgent.GetHandlerHistory:input_type -> cortex.axon.GetHandlerHistoryRequest - 13, // 39: cortex.axon.AxonAgent.Dispatch:input_type -> cortex.axon.DispatchRequest - 21, // 40: cortex.axon.AxonAgent.ReportInvocation:input_type -> cortex.axon.ReportInvocationRequest - 13, // 41: cortex.axon.AxonAgent.DispatchEvents:input_type -> cortex.axon.DispatchRequest - 18, // 42: cortex.axon.AxonAgent.RetryInvocation:input_type -> cortex.axon.RetryInvocationRequest - 7, // 43: cortex.axon.AxonAgent.RegisterHandler:output_type -> cortex.axon.RegisterHandlerResponse - 9, // 44: cortex.axon.AxonAgent.UnregisterHandler:output_type -> cortex.axon.UnregisterHandlerResponse - 12, // 45: cortex.axon.AxonAgent.ListHandlers:output_type -> cortex.axon.ListHandlersResponse - 26, // 46: cortex.axon.AxonAgent.GetHandlerHistory:output_type -> cortex.axon.GetHandlerHistoryResponse - 14, // 47: cortex.axon.AxonAgent.Dispatch:output_type -> cortex.axon.DispatchMessage - 22, // 48: cortex.axon.AxonAgent.ReportInvocation:output_type -> cortex.axon.ReportInvocationResponse - 17, // 49: cortex.axon.AxonAgent.DispatchEvents:output_type -> cortex.axon.DispatchEvent - 19, // 50: cortex.axon.AxonAgent.RetryInvocation:output_type -> cortex.axon.RetryInvocationResponse - 43, // [43:51] is the sub-list for method output_type - 35, // [35:43] is the sub-list for method input_type - 35, // [35:35] is the sub-list for extension type_name - 35, // [35:35] is the sub-list for extension extendee - 0, // [0:35] is the sub-list for field type_name + 29, // 13: cortex.axon.DispatchHandlerInvoke.deadline:type_name -> google.protobuf.Timestamp + 2, // 14: cortex.axon.InvocationFailure.reason:type_name -> cortex.axon.FailureReason + 3, // 15: cortex.axon.InvocationFailure.source:type_name -> cortex.axon.FailureSource + 15, // 16: cortex.axon.DispatchEvent.invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 16, // 17: cortex.axon.DispatchEvent.failure:type_name -> cortex.axon.InvocationFailure + 15, // 18: cortex.axon.RetryInvocationRequest.invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 28, // 19: cortex.axon.RetryInvocationResponse.error:type_name -> cortex.axon.Error + 29, // 20: cortex.axon.Log.timestamp:type_name -> google.protobuf.Timestamp + 15, // 21: cortex.axon.ReportInvocationRequest.handler_invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 29, // 22: cortex.axon.ReportInvocationRequest.start_client_timestamp:type_name -> google.protobuf.Timestamp + 23, // 23: cortex.axon.ReportInvocationRequest.result:type_name -> cortex.axon.InvokeResult + 28, // 24: cortex.axon.ReportInvocationRequest.error:type_name -> cortex.axon.Error + 20, // 25: cortex.axon.ReportInvocationRequest.logs:type_name -> cortex.axon.Log + 28, // 26: cortex.axon.ReportInvocationResponse.error:type_name -> cortex.axon.Error + 29, // 27: cortex.axon.GetHandlerHistoryRequest.start_time:type_name -> google.protobuf.Timestamp + 29, // 28: cortex.axon.GetHandlerHistoryRequest.end_time:type_name -> google.protobuf.Timestamp + 29, // 29: cortex.axon.HandlerExecution.publish_server_timestamp:type_name -> google.protobuf.Timestamp + 29, // 30: cortex.axon.HandlerExecution.receive_server_timestamp:type_name -> google.protobuf.Timestamp + 29, // 31: cortex.axon.HandlerExecution.start_client_timestamp:type_name -> google.protobuf.Timestamp + 28, // 32: cortex.axon.HandlerExecution.error:type_name -> cortex.axon.Error + 20, // 33: cortex.axon.HandlerExecution.logs:type_name -> cortex.axon.Log + 28, // 34: cortex.axon.GetHandlerHistoryResponse.error:type_name -> cortex.axon.Error + 25, // 35: cortex.axon.GetHandlerHistoryResponse.history:type_name -> cortex.axon.HandlerExecution + 4, // 36: cortex.axon.AxonAgent.RegisterHandler:input_type -> cortex.axon.RegisterHandlerRequest + 8, // 37: cortex.axon.AxonAgent.UnregisterHandler:input_type -> cortex.axon.UnregisterHandlerRequest + 10, // 38: cortex.axon.AxonAgent.ListHandlers:input_type -> cortex.axon.ListHandlersRequest + 24, // 39: cortex.axon.AxonAgent.GetHandlerHistory:input_type -> cortex.axon.GetHandlerHistoryRequest + 13, // 40: cortex.axon.AxonAgent.Dispatch:input_type -> cortex.axon.DispatchRequest + 21, // 41: cortex.axon.AxonAgent.ReportInvocation:input_type -> cortex.axon.ReportInvocationRequest + 13, // 42: cortex.axon.AxonAgent.DispatchEvents:input_type -> cortex.axon.DispatchRequest + 18, // 43: cortex.axon.AxonAgent.RetryInvocation:input_type -> cortex.axon.RetryInvocationRequest + 7, // 44: cortex.axon.AxonAgent.RegisterHandler:output_type -> cortex.axon.RegisterHandlerResponse + 9, // 45: cortex.axon.AxonAgent.UnregisterHandler:output_type -> cortex.axon.UnregisterHandlerResponse + 12, // 46: cortex.axon.AxonAgent.ListHandlers:output_type -> cortex.axon.ListHandlersResponse + 26, // 47: cortex.axon.AxonAgent.GetHandlerHistory:output_type -> cortex.axon.GetHandlerHistoryResponse + 14, // 48: cortex.axon.AxonAgent.Dispatch:output_type -> cortex.axon.DispatchMessage + 22, // 49: cortex.axon.AxonAgent.ReportInvocation:output_type -> cortex.axon.ReportInvocationResponse + 17, // 50: cortex.axon.AxonAgent.DispatchEvents:output_type -> cortex.axon.DispatchEvent + 19, // 51: cortex.axon.AxonAgent.RetryInvocation:output_type -> cortex.axon.RetryInvocationResponse + 44, // [44:52] is the sub-list for method output_type + 36, // [36:44] is the sub-list for method input_type + 36, // [36:36] is the sub-list for extension type_name + 36, // [36:36] is the sub-list for extension extendee + 0, // [0:36] is the sub-list for field type_name } func init() { file_cortex_axon_agent_proto_init() } diff --git a/sdks/go/README.md b/sdks/go/README.md index 96caadd4..0645eb25 100644 --- a/sdks/go/README.md +++ b/sdks/go/README.md @@ -108,3 +108,33 @@ This will begin executing your handler every second. + +## Failure hooks + +`WithPostDispatchHook` runs your own code after every invocation, and for +every failure, including ones only the agent sees (such as a webhook dropped +because every handler was busy). Use it to send failed messages somewhere you +can replay them from. + +```go +agent := axon.NewAxonAgent( + axon.WithPostDispatchHook(func(ctx context.Context, info axon.InvocationInfo, failure *axon.Failure) axon.PostAction { + if failure == nil { + return axon.PostActionIgnore // it worked + } + // info.Args holds the original webhook: "body", "content-type", "url" + deadLetters.Publish(info, failure) + if failure.Source == pb.FailureSource_RUNTIME && info.Attempt < 3 { + return axon.PostActionRetry + } + return axon.PostActionIgnore + }), +) +``` + +`failure.Reason` is `HANDLER_ERROR`, `TIMEOUT` or `NO_RESOURCES`. +`failure.Source` says whether handler code (`APPLICATION`) or Axon +(`RUNTIME`) failed. A timeout during a Cortex API call counts as `RUNTIME`. +`CallerRejected` means the sender was told the message wasn't accepted (a +webhook got a 503). `PostActionRetry` queues the invocation again as its next +attempt, up to 5 attempts. `WithPreDispatchHook` runs before each handler. diff --git a/sdks/go/axon_agent.go b/sdks/go/axon_agent.go index bd930f0f..d3524060 100644 --- a/sdks/go/axon_agent.go +++ b/sdks/go/axon_agent.go @@ -394,17 +394,27 @@ func (a *Agent) processDispatchStream(ctx context.Context, stream grpc.BidiStrea continue } + deadline, capped, expired := invokeDeadline(invoke, time.Now()) + if expired { + // The agent has already reported this one as failed; running + // it now would duplicate work a failure hook may be replaying. + a.logger.Warn("skipping invocation that expired before it could start", + zap.String("handler", invoke.HandlerName), + zap.String("invocation-id", invoke.InvocationId)) + continue + } + timeoutCtx := ctx cancel := func() {} - if invoke.TimeoutMs != 0 { - timeoutCtx, cancel = context.WithTimeout(context.Background(), time.Millisecond*time.Duration(invoke.TimeoutMs)) + if !deadline.IsZero() { + timeoutCtx, cancel = context.WithDeadline(context.Background(), deadline) } runningHandlers.Add(1) go func() { defer runningHandlers.Done() defer cancel() - a.invokeHandler(timeoutCtx, invoke) + a.invokeHandler(timeoutCtx, invoke, capped) }() case pb.DispatchMessageType_DISPATCH_MESSAGE_WORK_COMPLETED: a.logger.Info("work completed, shutting down") @@ -415,7 +425,35 @@ func (a *Agent) processDispatchStream(ctx context.Context, stream grpc.BidiStrea } } -func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInvoke) { +// deadlineMargin is how far ahead of the agent's deadline a handler must +// finish, so its report reaches the agent before the agent gives up on it. +const deadlineMargin = time.Second + +// invokeDeadline works out when a handler must finish: its own timeout, but no +// later than the agent's deadline. capped means the agent's deadline is the +// tighter one, i.e. the invocation waited before it got here. expired means +// it is already too late to start. +func invokeDeadline(invoke *pb.DispatchHandlerInvoke, now time.Time) (deadline time.Time, capped bool, expired bool) { + if invoke.TimeoutMs != 0 { + deadline = now.Add(time.Duration(invoke.TimeoutMs) * time.Millisecond) + } + if invoke.GetDeadline() == nil { + return deadline, false, false + } + agentDeadline := invoke.GetDeadline().AsTime().Add(-deadlineMargin) + if !now.Before(agentDeadline) { + return time.Time{}, false, true + } + if deadline.IsZero() || agentDeadline.Before(deadline) { + return agentDeadline, true, false + } + return deadline, false, false +} + +// invokeHandler runs one invocation. capped means its deadline came from the +// agent rather than its own timeout, so a timeout is time spent waiting, not +// the handler being slow. +func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInvoke, capped bool) { handlerInfo, ok := a.registeredHandlers[invoke.HandlerId] if !ok { a.logger.Error("handler not found", zap.String("handler", invoke.HandlerName)) @@ -485,6 +523,9 @@ func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInv if hc, ok := hctx.(*handlerContext); ok && hc.calls.blockedOnApi() { failure.Source = pb.FailureSource_RUNTIME failure.Message += " while waiting on a Cortex API call" + } else if capped { + failure.Source = pb.FailureSource_RUNTIME + failure.Message = "invocation waited for a free handler and ran out of time" } } diff --git a/sdks/go/hooks_test.go b/sdks/go/hooks_test.go index bddcbf7b..d052e616 100644 --- a/sdks/go/hooks_test.go +++ b/sdks/go/hooks_test.go @@ -13,6 +13,7 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/timestamppb" ) type hookCall struct { @@ -282,3 +283,30 @@ func TestDispatchRequestAdvertisesFailureEvents(t *testing.T) { }) } } + +func TestInvokeDeadline(t *testing.T) { + now := time.Now() + agentDeadline := func(d time.Duration) *timestamppb.Timestamp { return timestamppb.New(now.Add(d)) } + + // Older agent: only the timeout counts. + d, capped, expired := invokeDeadline(&pb.DispatchHandlerInvoke{TimeoutMs: 1000}, now) + require.Equal(t, now.Add(time.Second), d) + require.False(t, capped) + require.False(t, expired) + + // Plenty of time left: the handler's own timeout is tighter. + d, capped, expired = invokeDeadline(&pb.DispatchHandlerInvoke{TimeoutMs: 1000, Deadline: agentDeadline(time.Minute)}, now) + require.Equal(t, now.Add(time.Second), d) + require.False(t, capped) + require.False(t, expired) + + // It waited before arriving: the agent's deadline, less the margin, wins. + d, capped, expired = invokeDeadline(&pb.DispatchHandlerInvoke{TimeoutMs: 10000, Deadline: agentDeadline(3 * time.Second)}, now) + require.WithinDuration(t, now.Add(2*time.Second), d, time.Millisecond) + require.True(t, capped) + require.False(t, expired) + + // Too late to start: the agent has reported it already. + _, _, expired = invokeDeadline(&pb.DispatchHandlerInvoke{TimeoutMs: 10000, Deadline: agentDeadline(500 * time.Millisecond)}, now) + require.True(t, expired) +} diff --git a/sdks/python/README.md b/sdks/python/README.md index ba821ae3..fa6e20ee 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -106,3 +106,66 @@ client = AxonClient(scope=globals(), max_concurrent_handlers=16) ``` Set it to `1` for the old behaviour, where handlers run one at a time. + +## Failure hooks + +Pass `post_dispatch` to `AxonClient` to run your own code when an +invocation fails, for example to send the failed message somewhere you can +replay it from. The hook runs after every invocation (`failure` is `None` on +success). It also runs for failures that only the agent can see, such as a +webhook the agent couldn't hand to this process because every handler was +busy. Those failures arrive on their own stream, so they get through even +while every handler slot is taken. + +```python +from cortex_axon import AxonClient, FailureSource, PostAction +from cortex_axon.handler import cortex_webhook + + +@cortex_webhook(id="my-webhook-1") +def my_webhook_handler(context): + process(context.args["body"]) + + +def on_dispatch(info, failure): + if failure is None: + return # it worked + # info.args holds the original webhook: "body", "content-type", "url" + dead_letters.publish({ + "handler": info.handler_name, + "invocation_id": info.invocation_id, + "attempt": info.attempt, + "reason": failure.reason.name, # HANDLER_ERROR, TIMEOUT, NO_RESOURCES + "source": failure.source.name, # APPLICATION or RUNTIME + "caller_rejected": failure.caller_rejected, + "message": info.args, + }) + # Optionally have the agent run it again, as a new attempt. + if failure.source == FailureSource.RUNTIME and info.attempt < 3: + return PostAction.RETRY + + +client = AxonClient(scope=globals(), post_dispatch=on_dispatch) +client.run() +``` + +| Reason | When | Source | +|---|---|---| +| `HANDLER_ERROR` | The handler raised | `APPLICATION` | +| `TIMEOUT` | The handler ran past its timeout | `RUNTIME` if a Cortex API call was in flight or the agent never got a result; otherwise `APPLICATION` | +| `NO_RESOURCES` | The agent's queue was full, or the message expired waiting for a free handler | `RUNTIME` | + +`caller_rejected` is true when the sender was told the message wasn't +accepted (a webhook got a 503), so the sender may also retry it. + +Returning `PostAction.RETRY` asks the agent to queue the invocation again, +with `attempt` incremented. The agent stops after 5 attempts. Returning +nothing ignores the failure. + +`pre_dispatch` takes a function of `info` that runs before each handler, +e.g. to record every inbound message before processing it. + +Hooks run for every handler type, not only webhooks. Python can't stop a +handler that overruns its timeout: at the deadline the SDK reports the +`TIMEOUT` and runs the hook, and whatever the handler does after that is +dropped. diff --git a/sdks/python/cortex_axon/__init__.py b/sdks/python/cortex_axon/__init__.py index 8bafa392..1150881e 100644 --- a/sdks/python/cortex_axon/__init__.py +++ b/sdks/python/cortex_axon/__init__.py @@ -1,2 +1,7 @@ from .axon_client import AxonClient as AxonClient from .axon_client import CortexHTTPError as CortexHTTPError +from .hooks import Failure as Failure +from .hooks import FailureReason as FailureReason +from .hooks import FailureSource as FailureSource +from .hooks import InvocationInfo as InvocationInfo +from .hooks import PostAction as PostAction diff --git a/sdks/python/cortex_axon/axon_client.py b/sdks/python/cortex_axon/axon_client.py index e92c0650..e0c76735 100644 --- a/sdks/python/cortex_axon/axon_client.py +++ b/sdks/python/cortex_axon/axon_client.py @@ -5,7 +5,7 @@ import traceback from concurrent.futures import ThreadPoolExecutor from datetime import datetime -from time import monotonic, sleep +from time import monotonic, sleep, time from typing import Any, Optional from urllib.parse import quote_plus @@ -24,8 +24,37 @@ CortexAnnotation, find_annotated_methods, ) +from .hooks import ( + Failure, + FailureReason, + FailureSource, + InvocationInfo, + PostAction, + PostDispatchHook, + PreDispatchHook, +) DEFAULT_MAX_CONCURRENT_HANDLERS = 8 +RECONNECT_WAIT_SECONDS = 5 +# How far ahead of the agent's deadline a handler must finish, so its report +# reaches the agent before the agent gives up on it. +DEADLINE_MARGIN_SECONDS = 1.0 + + +def _invoke_timeout(invoke, now: float): + """How long a handler may run, in seconds (None for no limit): its own + timeout, but no later than the agent's deadline. capped means the agent's + deadline is the tighter one, i.e. the invocation waited before it got + here. expired means it is already too late to start.""" + timeout = invoke.timeout_ms / 1000 if invoke.timeout_ms > 0 else None + if not invoke.HasField("deadline"): + return timeout, False, False + remaining = invoke.deadline.ToMicroseconds() / 1e6 - now - DEADLINE_MARGIN_SECONDS + if remaining <= 0: + return None, False, True + if timeout is None or remaining < timeout: + return remaining, True, False + return timeout, False, False # TODO Add docstrings @@ -152,6 +181,8 @@ def __init__( scope=None, max_receive_message_size: Optional[int] = None, max_concurrent_handlers: int = DEFAULT_MAX_CONCURRENT_HANDLERS, + pre_dispatch: Optional[PreDispatchHook] = None, + post_dispatch: Optional[PostDispatchHook] = None, ): if max_concurrent_handlers < 1: raise ValueError("max_concurrent_handlers must be at least 1") @@ -196,6 +227,20 @@ def __init__( ) self._handler_slots = threading.BoundedSemaphore(max_concurrent_handlers) + self.pre_dispatch = pre_dispatch + self.post_dispatch = post_dispatch + # Failure events have their own thread and pool, apart from the + # handler slots: they arrive exactly when every handler is busy, and + # must never wait for one to free up. + self._hook_executor = None + self._events_thread = None + self._events_stop = threading.Event() + self._events_call = None + if post_dispatch: + self._hook_executor = ThreadPoolExecutor( + max_workers=2, thread_name_prefix="axon-hook" + ) + def _get_stub(self): if not self.agent_stub: self.agent_channel = grpc.insecure_channel( @@ -231,7 +276,98 @@ def _find_handler(self, name): return h return None - def _handler_context(self, args, timeout_ms: int = 0): + def _run_pre_dispatch(self, info: InvocationInfo): + if not self.pre_dispatch: + return + try: + self.pre_dispatch(info) + except Exception: + print(f"Error in pre_dispatch hook for {info.handler_name} (id={info.invocation_id})") + traceback.print_exc() + + def _run_post_dispatch(self, info: InvocationInfo, failure: Optional[Failure], stub): + if not self.post_dispatch: + return + try: + action = self.post_dispatch(info, failure) + except Exception: + print(f"Error in post_dispatch hook for {info.handler_name} (id={info.invocation_id})") + traceback.print_exc() + return + if action == PostAction.RETRY and failure is not None: + self._retry_invocation(info, stub) + + def _retry_invocation(self, info: InvocationInfo, stub): + try: + response = stub.RetryInvocation( + cortex_axon_agent_pb2.RetryInvocationRequest(invoke=info.invoke) + ) + except grpc.RpcError as e: + if e.code() == grpc.StatusCode.UNIMPLEMENTED: + print("Agent does not support retries, upgrade the agent") + else: + print(f"Error retrying invocation {info.invocation_id}: {e}") + return + if response.error and response.error.code: + print(f"Agent rejected retry of {info.invocation_id}: {response.error.message}") + return + print( + f"Retrying {info.handler_name} (id={info.invocation_id}) " + f"as id={response.invocation_id} attempt={info.attempt + 1}" + ) + + def _start_failure_events(self): + if not self.post_dispatch or self._events_thread: + return + self._events_thread = threading.Thread( + target=self._run_failure_events, name="axon-failure-events", daemon=True + ) + self._events_thread.start() + + def _stop_failure_events(self): + self._events_stop.set() + call = self._events_call + if call: + call.cancel() + + def _run_failure_events(self): + # A channel of its own, so the dispatch loop resetting its channel + # after an error doesn't cut this stream too. + channel = grpc.insecure_channel( + self.agent_hostport, options=self.grpc_channel_options + ) + stub = cortex_axon_agent_pb2_grpc.AxonAgentStub(channel) + try: + while not self._events_stop.is_set(): + try: + self._events_call = stub.DispatchEvents( + cortex_axon_agent_pb2.DispatchRequest( + dispatch_id=self.id, + client_version=AGENT_VERSION, + failure_events=True, + ) + ) + for event in self._events_call: + if not event.HasField("invoke"): + continue + self._hook_executor.submit( + self._run_post_dispatch, + InvocationInfo.from_proto(event.invoke), + Failure.from_proto(event.failure), + stub, + ) + except grpc.RpcError as e: + if self._events_stop.is_set(): + return + if e.code() == grpc.StatusCode.UNIMPLEMENTED: + print("Agent does not send failure events, upgrade the agent to receive them") + return + print(f"Failure event stream ended ({e.code()}), reconnecting in {RECONNECT_WAIT_SECONDS} seconds...") + self._events_stop.wait(RECONNECT_WAIT_SECONDS) + finally: + channel.close() + + def _handler_context(self, args, timeout_ms: float = 0): return HandlerContext(self.cortex_stub, args, timeout_ms) def _submit_handler(self, handler_to_invoke, invoke): @@ -239,6 +375,16 @@ def _submit_handler(self, handler_to_invoke, invoke): # sees its own queue fill rather than this process growing without # bound, and the dispatch stream stays readable up to that point. self._handler_slots.acquire() + _, _, expired = _invoke_timeout(invoke, time()) + if expired: + # The agent has already reported this one as failed; running it + # now would duplicate work a failure hook may be replaying. + self._handler_slots.release() + print( + f"Dispatch handler SKIPPED {invoke.handler_name} (id={invoke.invocation_id}): " + "expired while waiting for a free handler" + ) + return stub = self.agent_stub try: future = self._executor.submit( @@ -255,11 +401,58 @@ def _invoke_handler(self, handler_to_invoke, invoke, stub): ) handler_err = None handler_result = None + failure = None result = None + info = InvocationInfo.from_proto(invoke) + self._run_pre_dispatch(info) now = datetime.now() - ctx = self._handler_context(invoke.args, invoke.timeout_ms) + timeout, capped, _ = _invoke_timeout(invoke, time()) + ctx = self._handler_context(invoke.args, timeout * 1000 if timeout else 0) duration_ms = 0 + # A Python thread can't be stopped, so at the deadline the timer + # reports the timeout itself and the handler's eventual result is + # dropped. Whichever finishes first claims the invocation. + finish_lock = threading.Lock() + finished = False + + def claim(): + nonlocal finished + with finish_lock: + if finished: + return False + finished = True + return True + + def on_timeout(): + if not claim(): + return + source = FailureSource.APPLICATION + message = f"handler did not finish within {invoke.timeout_ms}ms" + if ctx.blocked_on_api(): + source = FailureSource.RUNTIME + message += " while waiting on a Cortex API call" + elif capped: + source = FailureSource.RUNTIME + message = "invocation waited for a free handler and ran out of time" + print(f"Dispatch handler TIMEOUT {invoke.handler_name} (id={invoke.invocation_id}): {message}") + self._report(stub, cortex_axon_agent_pb2.ReportInvocationRequest( + handler_invoke=invoke, + start_client_timestamp=now, + duration_ms=invoke.timeout_ms, + error=common_pb2.Error(code="timeout", message=message), + logs=list(ctx.logs), + )) + self._run_post_dispatch( + info, Failure(reason=FailureReason.TIMEOUT, source=source, message=message), stub + ) + + timer = None + if timeout: + timer = threading.Timer(timeout, on_timeout) + timer.daemon = True + timer.start() + try: result = handler_to_invoke(ctx) duration_ms = int( @@ -276,6 +469,24 @@ def _invoke_handler(self, handler_to_invoke, invoke, stub): handler_err = common_pb2.Error( code="unexpected", message="Error calling handler: " + str(e) ) + if ctx.blocked_on_api(): + failure = Failure( + reason=FailureReason.TIMEOUT, + source=FailureSource.RUNTIME, + message=f"Cortex API call did not finish within the {invoke.timeout_ms}ms timeout: {e}", + ) + else: + failure = Failure( + reason=FailureReason.HANDLER_ERROR, + source=FailureSource.APPLICATION, + message=str(e), + ) + + if timer: + timer.cancel() + if not claim(): + print(f"Dispatch handler {invoke.handler_name} (id={invoke.invocation_id}) finished after its timeout, result dropped") + return invoke_info = cortex_axon_agent_pb2.ReportInvocationRequest( handler_invoke=invoke, @@ -286,24 +497,36 @@ def _invoke_handler(self, handler_to_invoke, invoke, stub): logs=ctx.logs, ) + self._report(stub, invoke_info) + + # After the report, so a slow hook can't make the agent decide this + # invocation was never answered and report the failure a second time. + self._run_post_dispatch(info, failure, stub) + + def _report(self, stub, invoke_info): # A failure here must not escape the worker: the dispatch loop is no # longer in the call path, so an exception would be swallowed by the # future and the invocation would look like it vanished. try: report_response = stub.ReportInvocation(invoke_info) + if report_response.error and report_response.error.code: + print( + f"Error reporting invocation: {report_response.error}" + ) except Exception as e: print( - f"Error reporting invocation {invoke.invocation_id}: {e}" + f"Error reporting invocation {invoke_info.handler_invoke.invocation_id}: {e}" ) traceback.print_exc() - return - - if report_response.error and report_response.error.code: - print( - f"Error reporting invocation: {report_response.error}" - ) def run(self): + self._start_failure_events() + try: + self._run_dispatch() + finally: + self._stop_failure_events() + + def _run_dispatch(self): registered = False while True: stub = self._get_stub() @@ -314,7 +537,8 @@ def run(self): initial_message = cortex_axon_agent_pb2.DispatchRequest( dispatch_id=self.id, - client_version=AGENT_VERSION + client_version=AGENT_VERSION, + failure_events=self.post_dispatch is not None, ) for response in stub.Dispatch( _client_message_iterator(initial_message) @@ -364,6 +588,14 @@ def __init__(self, api: cortex_api_pb2_grpc.CortexApiStub, args: dict, timeout_m # worker after the agent has given up on the invocation, and once # every worker is held the client stops taking new invocations. self.deadline = monotonic() + timeout_ms / 1000 if timeout_ms > 0 else None + # Whether a Cortex API call was running, or ran out of time, when the + # deadline hit, which makes the timeout Axon's fault rather than the + # handler's. + self._api_in_flight = 0 + self._api_timed_out = False + + def blocked_on_api(self) -> bool: + return self._api_in_flight > 0 or self._api_timed_out def log(self, message: str, level="INFO"): """ @@ -395,5 +627,15 @@ def cortex_api_call( timeout = None if self.deadline is not None: timeout = max(self.deadline - monotonic(), 0) - response = self.api.Call(request, timeout=timeout) + self._api_in_flight += 1 + try: + response = self.api.Call(request, timeout=timeout) + except grpc.RpcError as e: + if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED: + # Set before _api_in_flight drops, so a timeout check never + # sees neither. + self._api_timed_out = True + raise + finally: + self._api_in_flight -= 1 return CortexResponse(response.status_code, response.body, response.headers) diff --git a/sdks/python/cortex_axon/hooks.py b/sdks/python/cortex_axon/hooks.py new file mode 100644 index 00000000..6b910c54 --- /dev/null +++ b/sdks/python/cortex_axon/hooks.py @@ -0,0 +1,85 @@ +from dataclasses import dataclass, field +from enum import Enum, IntEnum +from typing import Any, Callable, Optional + +from generated import cortex_axon_agent_pb2 + + +class FailureReason(IntEnum): + UNSPECIFIED = cortex_axon_agent_pb2.FAILURE_REASON_UNSPECIFIED + # The handler raised. + HANDLER_ERROR = cortex_axon_agent_pb2.HANDLER_ERROR + # The invocation did not finish within its timeout. + TIMEOUT = cortex_axon_agent_pb2.TIMEOUT + # The invocation never ran: the agent's queue was full, or it expired + # waiting for a free handler. + NO_RESOURCES = cortex_axon_agent_pb2.NO_RESOURCES + + +class FailureSource(IntEnum): + UNSPECIFIED = cortex_axon_agent_pb2.FAILURE_SOURCE_UNSPECIFIED + # Handler code failed. + APPLICATION = cortex_axon_agent_pb2.APPLICATION + # Axon failed: the agent, the SDK, or a Cortex API call. + RUNTIME = cortex_axon_agent_pb2.RUNTIME + + +class PostAction(Enum): + # Leave the failure as it is. + IGNORE = "ignore" + # Queue the invocation again as its next attempt. + RETRY = "retry" + + +@dataclass(frozen=True) +class InvocationInfo: + """Identifies one invocation. For a webhook, args holds the inbound + message: "body", "content-type" and "url".""" + + invocation_id: str + handler_name: str + invoke_type: int + args: dict[str, str] + # 1 for the first delivery, counting up with each retry. + attempt: int + invoke: Any = field(repr=False, compare=False) + + @classmethod + def from_proto(cls, invoke: cortex_axon_agent_pb2.DispatchHandlerInvoke) -> "InvocationInfo": + return cls( + invocation_id=invoke.invocation_id, + handler_name=invoke.handler_name, + invoke_type=invoke.reason, + args=dict(invoke.args), + attempt=max(invoke.attempt, 1), + invoke=invoke, + ) + + +@dataclass(frozen=True) +class Failure: + reason: FailureReason + source: FailureSource + message: str + # Whoever sent the invocation was told it was not accepted (a webhook + # sender got a 503), so it may arrive again. + caller_rejected: bool = False + + @classmethod + def from_proto(cls, failure: cortex_axon_agent_pb2.InvocationFailure) -> "Failure": + return cls( + reason=FailureReason(failure.reason), + source=FailureSource(failure.source), + message=failure.message, + caller_rejected=failure.caller_rejected, + ) + + +# Runs before each handler. +PreDispatchHook = Callable[[InvocationInfo], None] + +# Runs after each invocation with failure=None on success, and for every +# failure the agent detects on its own, such as an invocation that never +# reached this process. Returning None means PostAction.IGNORE; the return +# value only matters for a failure. +PostDispatchHook = Callable[[InvocationInfo, Optional[Failure]], Optional[PostAction]] diff --git a/sdks/python/example/main.py b/sdks/python/example/main.py index 826e7b4c..ca92ffe4 100644 --- a/sdks/python/example/main.py +++ b/sdks/python/example/main.py @@ -1,5 +1,6 @@ import json +from cortex_axon import PostAction from cortex_axon.axon_client import AxonClient, HandlerContext from cortex_axon.handler import cortex_handler, cortex_scheduled @@ -50,10 +51,25 @@ def my_invoke_handler(context: HandlerContext) -> str: return json.dumps(result) +def on_dispatch(info, failure): + # Runs after every invocation, and for failures only the agent sees. + if failure is None: + return None + print( + f"FAILED {info.handler_name} (id={info.invocation_id}, attempt={info.attempt}) " + f"reason={failure.reason.name} source={failure.source.name}: {failure.message}" + ) + # A real hook would save info.args somewhere it can be replayed from. + if info.attempt < 3: + return PostAction.RETRY + return PostAction.IGNORE + + def run(): # Connect to the gRPC server client = AxonClient( scope=globals(), + post_dispatch=on_dispatch, ) client.run() diff --git a/sdks/python/tests/test_hooks.py b/sdks/python/tests/test_hooks.py new file mode 100644 index 00000000..360e7953 --- /dev/null +++ b/sdks/python/tests/test_hooks.py @@ -0,0 +1,370 @@ +import contextlib +import queue +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import grpc +import pytest + +from cortex_axon import ( + AxonClient, + FailureReason, + FailureSource, + PostAction, +) +from cortex_axon.handler import CortexWebhook +from generated import ( + common_pb2, + cortex_api_pb2_grpc, + cortex_axon_agent_pb2, + cortex_axon_agent_pb2_grpc, +) + +HANDLER = "my_hooked_webhook" + + +class HookAgent(cortex_axon_agent_pb2_grpc.AxonAgentServicer): + """Sends a fixed set of invocations and failure events, then holds both + streams open, recording what the client sends back.""" + + def __init__(self, invokes=(), events=(), events_code=None): + self.invokes = list(invokes) + self.events = list(events) + self.events_code = events_code + self.dispatch_requests = queue.Queue() + self.retries = queue.Queue() + self.events_calls = 0 + self.reports = queue.Queue() + self.stop = threading.Event() + + def RegisterHandler(self, request, context): + return cortex_axon_agent_pb2.RegisterHandlerResponse( + error=common_pb2.Error(code="0", message="OK"), + id=request.dispatch_id, + ) + + def Dispatch(self, request_iterator, context): + request = next(request_iterator) + self.dispatch_requests.put(request) + for invoke in self.invokes: + yield cortex_axon_agent_pb2.DispatchMessage( + type=cortex_axon_agent_pb2.DISPATCH_MESSAGE_INVOKE, + invoke=invoke, + ) + self.invokes = [] + self.stop.wait(timeout=10) + + def DispatchEvents(self, request, context): + self.events_calls += 1 + if self.events_code: + context.abort(self.events_code, "not here") + events, self.events = self.events, [] + yield from events + self.stop.wait(timeout=10) + + def ReportInvocation(self, request, context): + self.reports.put(request) + return cortex_axon_agent_pb2.ReportInvocationResponse() + + def RetryInvocation(self, request, context): + self.retries.put(request) + return cortex_axon_agent_pb2.RetryInvocationResponse(invocation_id="retry-1") + + +class SlowApi(cortex_api_pb2_grpc.CortexApiServicer): + def __init__(self): + self.release = threading.Event() + + def Call(self, request, context): + self.release.wait(timeout=10) + context.abort(grpc.StatusCode.UNAVAILABLE, "too late") + + +def invoke(invocation_id="1", timeout_ms=0, deadline_in=None, **args): + msg = cortex_axon_agent_pb2.DispatchHandlerInvoke( + invocation_id=invocation_id, + handler_id=HANDLER, + handler_name=HANDLER, + reason=cortex_axon_agent_pb2.WEBHOOK, + timeout_ms=timeout_ms, + args=args, + ) + if deadline_in is not None: + msg.deadline.FromMicroseconds(int((time.time() + deadline_in) * 1e6)) + return msg + + +def serve(servicer, add): + server = grpc.server(ThreadPoolExecutor(max_workers=8)) + port = server.add_insecure_port("127.0.0.1:0") + add(servicer, server) + server.start() + return server, port + + +class Hooks: + def __init__(self, action=None): + self.action = action + self.pre = queue.Queue() + self.post = queue.Queue() + + def pre_dispatch(self, info): + self.pre.put(info) + + def post_dispatch(self, info, failure): + self.post.put((info, failure)) + return self.action + + +@pytest.fixture +def run_client(): + servers = [] + agents = [] + clients = [] + + def start(agent, handler, hooks=None, api=None, max_concurrent_handlers=4): + server, port = serve(agent, cortex_axon_agent_pb2_grpc.add_AxonAgentServicer_to_server) + servers.append(server) + agents.append(agent) + api_port = port + if api: + api_server, api_port = serve(api, cortex_api_pb2_grpc.add_CortexApiServicer_to_server) + servers.append(api_server) + + webhook = CortexWebhook(id="hook") + webhook.set_func(handler) + client = AxonClient( + agent_host="127.0.0.1", + agent_port=port, + cortex_host="127.0.0.1", + cortex_port=api_port, + handlers=[webhook], + max_concurrent_handlers=max_concurrent_handlers, + pre_dispatch=hooks.pre_dispatch if hooks else None, + post_dispatch=hooks.post_dispatch if hooks else None, + ) + threading.Thread(target=client.run, daemon=True).start() + clients.append(client) + return client + + yield start + for agent in agents: + agent.stop.set() + for server in servers: + server.stop(grace=None) + # Drain handler and hook threads so none is still printing when the + # interpreter exits. + for client in clients: + client._stop_failure_events() + client._executor.shutdown(wait=True, cancel_futures=True) + if client._hook_executor: + client._hook_executor.shutdown(wait=True, cancel_futures=True) + if client._events_thread: + client._events_thread.join(timeout=5) + + +def test_post_dispatch_on_success(run_client): + hooks = Hooks(action=PostAction.RETRY) + agent = HookAgent(invokes=[invoke(body="payload")]) + + def my_hooked_webhook(ctx): + return "ok" + + run_client(agent, my_hooked_webhook, hooks) + + pre = hooks.pre.get(timeout=5) + info, failure = hooks.post.get(timeout=5) + assert pre.invocation_id == info.invocation_id == "1" + assert failure is None + assert info.args == {"body": "payload"} + assert info.attempt == 1 + # RETRY on success is ignored. + with pytest.raises(queue.Empty): + agent.retries.get(timeout=0.5) + + +def test_post_dispatch_on_handler_error_and_retry(run_client): + hooks = Hooks(action=PostAction.RETRY) + agent = HookAgent(invokes=[invoke(body="payload")]) + + def my_hooked_webhook(ctx): + raise ValueError("ohnoes") + + run_client(agent, my_hooked_webhook, hooks) + + info, failure = hooks.post.get(timeout=5) + assert failure.reason == FailureReason.HANDLER_ERROR + assert failure.source == FailureSource.APPLICATION + assert failure.message == "ohnoes" + + retry = agent.retries.get(timeout=5) + assert retry.invoke.invocation_id == "1" + assert retry.invoke.args["body"] == "payload" + + +def test_api_timeout_is_runtime(run_client): + hooks = Hooks() + api = SlowApi() + agent = HookAgent(invokes=[invoke(timeout_ms=200)]) + + def my_hooked_webhook(ctx): + ctx.cortex_api_call("/api/v1/slow", "GET") + + try: + run_client(agent, my_hooked_webhook, hooks, api=api) + _, failure = hooks.post.get(timeout=5) + finally: + api.release.set() + assert failure.reason == FailureReason.TIMEOUT + assert failure.source == FailureSource.RUNTIME + + +def test_slow_handler_times_out_as_application(run_client): + hooks = Hooks() + agent = HookAgent(invokes=[invoke(timeout_ms=200)]) + finished = threading.Event() + + def my_hooked_webhook(ctx): + time.sleep(0.6) + finished.set() + + run_client(agent, my_hooked_webhook, hooks) + + _, failure = hooks.post.get(timeout=5) + assert failure.reason == FailureReason.TIMEOUT + assert failure.source == FailureSource.APPLICATION + report = agent.reports.get(timeout=5) + assert report.error.code == "timeout" + + # The late finish is dropped: no second report, no second hook call. + assert finished.wait(timeout=5) + with pytest.raises(queue.Empty): + hooks.post.get(timeout=0.5) + assert agent.reports.empty() + + +def test_expired_invocation_is_not_started(run_client): + """The agent has already reported an invocation past its deadline, so the + SDK must not run it: a failure hook may be replaying it.""" + hooks = Hooks() + ran = queue.Queue() + agent = HookAgent(invokes=[ + invoke("stale", timeout_ms=10000, deadline_in=0.5), + invoke("fresh", timeout_ms=10000, deadline_in=60), + ]) + + def my_hooked_webhook(ctx): + ran.put(ctx.args) + + run_client(agent, my_hooked_webhook, hooks) + + info, failure = hooks.post.get(timeout=5) + assert info.invocation_id == "fresh" + assert failure is None + assert agent.reports.get(timeout=5).handler_invoke.invocation_id == "fresh" + with pytest.raises(queue.Empty): + hooks.post.get(timeout=0.5) + assert ran.qsize() == 1 + + +def test_timeout_capped_by_agent_deadline_is_runtime(run_client): + hooks = Hooks() + # Its own timeout is long, but the agent gives up in 1.5s: it waited. + agent = HookAgent(invokes=[invoke(timeout_ms=60000, deadline_in=1.5)]) + + def my_hooked_webhook(ctx): + time.sleep(2) + + run_client(agent, my_hooked_webhook, hooks) + + _, failure = hooks.post.get(timeout=5) + assert failure.reason == FailureReason.TIMEOUT + assert failure.source == FailureSource.RUNTIME + assert agent.reports.get(timeout=5).error.code == "timeout" + + +def test_hook_errors_are_contained(run_client): + agent = HookAgent(invokes=[invoke("1"), invoke("2")]) + ran = queue.Queue() + + class Broken(Hooks): + def post_dispatch(self, info, failure): + raise RuntimeError("hook bug") + + def my_hooked_webhook(ctx): + ran.put(True) + + run_client(agent, my_hooked_webhook, Broken()) + ran.get(timeout=5) + ran.get(timeout=5) + + +def test_dispatch_request_advertises_failure_events(run_client): + with_hooks = HookAgent() + run_client(with_hooks, lambda ctx: None, Hooks()) + assert with_hooks.dispatch_requests.get(timeout=5).failure_events + + without_hooks = HookAgent() + run_client(without_hooks, lambda ctx: None) + assert not without_hooks.dispatch_requests.get(timeout=5).failure_events + assert without_hooks.events_calls == 0 + + +def test_failure_events_reach_hook_while_handlers_are_saturated(run_client): + """The agent reports lost invocations exactly when every handler slot is + busy. Those events must not wait behind the slots.""" + slots = 2 + release = threading.Event() + started = threading.Barrier(slots + 1) + events = [ + cortex_axon_agent_pb2.DispatchEvent( + invoke=invoke(f"lost-{i}", body=f"msg-{i}"), + failure=cortex_axon_agent_pb2.InvocationFailure( + reason=cortex_axon_agent_pb2.NO_RESOURCES, + source=cortex_axon_agent_pb2.RUNTIME, + message="dispatch queue is full", + caller_rejected=True, + ), + ) + for i in range(20) + ] + # More invocations than slots, so the dispatch reader blocks. + agent = HookAgent(invokes=[invoke(str(i)) for i in range(slots + 2)], events=events) + hooks = Hooks(action=PostAction.RETRY) + + def my_hooked_webhook(ctx): + with contextlib.suppress(threading.BrokenBarrierError): + started.wait(timeout=5) + release.wait(timeout=10) + + try: + run_client(agent, my_hooked_webhook, hooks, max_concurrent_handlers=slots) + started.wait(timeout=5) + + seen = [hooks.post.get(timeout=5) for _ in events] + bodies = sorted(info.args["body"] for info, _ in seen) + assert bodies == sorted(f"msg-{i}" for i in range(20)) + for _, failure in seen: + assert failure.reason == FailureReason.NO_RESOURCES + assert failure.source == FailureSource.RUNTIME + assert failure.caller_rejected + retried = {agent.retries.get(timeout=5).invoke.invocation_id for _ in events} + assert retried == {f"lost-{i}" for i in range(20)} + finally: + release.set() + started.abort() + + +def test_older_agent_without_failure_events(run_client): + ran = queue.Queue() + agent = HookAgent(invokes=[invoke()], events_code=grpc.StatusCode.UNIMPLEMENTED) + + def my_hooked_webhook(ctx): + ran.put(True) + + client = run_client(agent, my_hooked_webhook, Hooks()) + ran.get(timeout=5) + client._events_thread.join(timeout=5) + assert not client._events_thread.is_alive() + assert agent.events_calls == 1 From a214e76bd07773e0a19101158a760f1cd1bd99a3 Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Wed, 11 Nov 2026 10:59:00 +1100 Subject: [PATCH 3/5] SDKs: replace pre/post-dispatch hooks with a single failure hook The post-dispatch hook ran on success too, so every hook began by checking for a failure, and the pre-dispatch hook only made sense paired with that success call. Neither was part of the ask. Python now takes on_failure and Go WithFailureHook; both run only for failures and return FailureAction (IGNORE / RETRY). Co-Authored-By: Claude Opus 5.5 --- sdks/go/README.md | 20 +++--- sdks/go/axon_agent.go | 17 ++--- sdks/go/hooks.go | 67 ++++++------------- sdks/go/hooks_test.go | 92 +++++++++++--------------- sdks/go/options.go | 3 +- sdks/python/README.md | 32 ++++----- sdks/python/cortex_axon/__init__.py | 2 +- sdks/python/cortex_axon/axon_client.py | 44 +++++------- sdks/python/cortex_axon/hooks.py | 14 ++-- sdks/python/example/main.py | 14 ++-- sdks/python/tests/test_hooks.py | 64 +++++++++--------- 11 files changed, 146 insertions(+), 223 deletions(-) diff --git a/sdks/go/README.md b/sdks/go/README.md index 0645eb25..4dec063c 100644 --- a/sdks/go/README.md +++ b/sdks/go/README.md @@ -111,23 +111,19 @@ This will begin executing your handler every second. ## Failure hooks -`WithPostDispatchHook` runs your own code after every invocation, and for -every failure, including ones only the agent sees (such as a webhook dropped -because every handler was busy). Use it to send failed messages somewhere you -can replay them from. +`WithFailureHook` runs your own code for every failed invocation, including +ones only the agent sees (such as a webhook dropped because every handler was +busy). Use it to send failed messages somewhere you can replay them from. ```go agent := axon.NewAxonAgent( - axon.WithPostDispatchHook(func(ctx context.Context, info axon.InvocationInfo, failure *axon.Failure) axon.PostAction { - if failure == nil { - return axon.PostActionIgnore // it worked - } + axon.WithFailureHook(func(ctx context.Context, info axon.InvocationInfo, failure *axon.Failure) axon.FailureAction { // info.Args holds the original webhook: "body", "content-type", "url" deadLetters.Publish(info, failure) if failure.Source == pb.FailureSource_RUNTIME && info.Attempt < 3 { - return axon.PostActionRetry + return axon.FailureActionRetry } - return axon.PostActionIgnore + return axon.FailureActionIgnore }), ) ``` @@ -136,5 +132,5 @@ agent := axon.NewAxonAgent( `failure.Source` says whether handler code (`APPLICATION`) or Axon (`RUNTIME`) failed. A timeout during a Cortex API call counts as `RUNTIME`. `CallerRejected` means the sender was told the message wasn't accepted (a -webhook got a 503). `PostActionRetry` queues the invocation again as its next -attempt, up to 5 attempts. `WithPreDispatchHook` runs before each handler. +webhook got a 503). `FailureActionRetry` queues the invocation again as its +next attempt, up to 5 attempts. diff --git a/sdks/go/axon_agent.go b/sdks/go/axon_agent.go index d3524060..7c0cbd96 100644 --- a/sdks/go/axon_agent.go +++ b/sdks/go/axon_agent.go @@ -37,8 +37,7 @@ type Agent struct { sleepOnError time.Duration done chan struct{} - preDispatch PreDispatchHook - postDispatch PostDispatchHook + onFailure FailureHook } // NewAxonAgent creates a new AxonAgent with the specified options. You @@ -61,8 +60,7 @@ func NewAxonAgent(options ...Option) *Agent { logger: logger, sleepOnError: ao.sleepOnError, done: make(chan struct{}), - preDispatch: ao.preDispatch, - postDispatch: ao.postDispatch, + onFailure: ao.onFailure, } a.logger = logger @@ -256,7 +254,7 @@ func (a *Agent) Run(ctx context.Context) error { // retries from a handler still has a live agent connection. eventsCtx, stopEvents := context.WithCancel(context.Background()) eventsDone := make(chan struct{}) - if a.postDispatch != nil { + if a.onFailure != nil { go func() { defer close(eventsDone) a.runFailureEvents(eventsCtx) @@ -309,7 +307,7 @@ func (a *Agent) Run(ctx context.Context) error { err = stream.Send(&pb.DispatchRequest{ DispatchId: a.DispatchId, ClientVersion: version.Client, - FailureEvents: a.postDispatch != nil, + FailureEvents: a.onFailure != nil, }) if err != nil { a.logger.Error("failed to send registration id", zap.Error(err)) @@ -460,9 +458,6 @@ func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInv return } - info := newInvocationInfo(invoke) - a.runPreDispatch(info) - report := &pb.ReportInvocationRequest{ HandlerInvoke: invoke, StartClientTimestamp: timestamppb.Now(), @@ -555,7 +550,9 @@ func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInv // After the report, so a slow hook can't make the agent decide this // invocation was never answered and report the failure a second time. - a.runPostDispatch(info, failure) + if failure != nil { + a.runFailureHook(newInvocationInfo(invoke), failure) + } } func (a *Agent) executeHandlerWithRecover(handler *handlerInfo, ctx HandlerContext) (result any, d time.Duration, err error) { diff --git a/sdks/go/hooks.go b/sdks/go/hooks.go index c3e8234c..f3faf7ee 100644 --- a/sdks/go/hooks.go +++ b/sdks/go/hooks.go @@ -13,7 +13,7 @@ import ( "google.golang.org/grpc/status" ) -// InvocationInfo identifies one invocation for the dispatch hooks. For a +// InvocationInfo identifies one invocation for the failure hook. For a // webhook, Args holds the inbound message: "body", "content-type" and "url". type InvocationInfo struct { InvocationID string @@ -57,69 +57,44 @@ func failureFromProto(f *pb.InvocationFailure) *Failure { } } -// PostAction is what a PostDispatchHook wants done with a failed invocation. -type PostAction int +// FailureAction is what a FailureHook wants done with a failed invocation. +type FailureAction int const ( - // PostActionIgnore leaves the failure as it is. - PostActionIgnore PostAction = iota - // PostActionRetry queues the invocation again as its next attempt. - PostActionRetry + // FailureActionIgnore leaves the failure as it is. + FailureActionIgnore FailureAction = iota + // FailureActionRetry queues the invocation again as its next attempt. + FailureActionRetry ) -// PreDispatchHook runs before each invocation's handler. -type PreDispatchHook func(ctx context.Context, info InvocationInfo) +// FailureHook runs for every failed invocation, including failures the agent +// detects on its own, such as an invocation that never reached this process. +type FailureHook func(ctx context.Context, info InvocationInfo, failure *Failure) FailureAction -// PostDispatchHook runs after each invocation, with a nil failure when it -// succeeded. It also runs for failures the agent detects on its own, such as -// an invocation that never reached this process, so it sees every failed -// message. Its return value only matters for a failure. -type PostDispatchHook func(ctx context.Context, info InvocationInfo, failure *Failure) PostAction - -// WithPreDispatchHook sets a hook that runs before every handler. -func WithPreDispatchHook(hook PreDispatchHook) Option { +// WithFailureHook sets a hook that runs for every failed invocation. +func WithFailureHook(hook FailureHook) Option { return func(a *agentOptions) { - a.preDispatch = hook + a.onFailure = hook } } -// WithPostDispatchHook sets a hook that runs after every invocation and for -// every failure, including ones the agent detects. -func WithPostDispatchHook(hook PostDispatchHook) Option { - return func(a *agentOptions) { - a.postDispatch = hook - } -} - -func (a *Agent) runPreDispatch(info InvocationInfo) { - if a.preDispatch == nil { - return - } - defer func() { - if r := recover(); r != nil { - a.logger.Error("panic in pre-dispatch hook", zap.Any("panic", r)) - } - }() - a.preDispatch(context.Background(), info) -} - -func (a *Agent) runPostDispatch(info InvocationInfo, failure *Failure) { - if a.postDispatch == nil { +func (a *Agent) runFailureHook(info InvocationInfo, failure *Failure) { + if a.onFailure == nil { return } - if a.callPostDispatch(info, failure) == PostActionRetry && failure != nil { + if a.callFailureHook(info, failure) == FailureActionRetry { a.retryInvocation(info) } } -func (a *Agent) callPostDispatch(info InvocationInfo, failure *Failure) (action PostAction) { +func (a *Agent) callFailureHook(info InvocationInfo, failure *Failure) (action FailureAction) { defer func() { if r := recover(); r != nil { - a.logger.Error("panic in post-dispatch hook", zap.Any("panic", r)) - action = PostActionIgnore + a.logger.Error("panic in failure hook", zap.Any("panic", r)) + action = FailureActionIgnore } }() - return a.postDispatch(context.Background(), info, failure) + return a.onFailure(context.Background(), info, failure) } func (a *Agent) retryInvocation(info InvocationInfo) { @@ -198,6 +173,6 @@ func (a *Agent) readFailureEvents(ctx context.Context) error { if ev.GetInvoke() == nil { continue } - a.runPostDispatch(newInvocationInfo(ev.GetInvoke()), failureFromProto(ev.GetFailure())) + a.runFailureHook(newInvocationInfo(ev.GetInvoke()), failureFromProto(ev.GetFailure())) } } diff --git a/sdks/go/hooks_test.go b/sdks/go/hooks_test.go index d052e616..0b3d1548 100644 --- a/sdks/go/hooks_test.go +++ b/sdks/go/hooks_test.go @@ -21,19 +21,19 @@ type hookCall struct { failure *Failure } -// hookRecorder is a post-dispatch hook that records each call. +// hookRecorder is a failure hook that records each call. type hookRecorder struct { mu sync.Mutex calls []hookCall - action PostAction + action FailureAction seen chan struct{} } -func newHookRecorder(action PostAction) *hookRecorder { +func newHookRecorder(action FailureAction) *hookRecorder { return &hookRecorder{action: action, seen: make(chan struct{}, 10)} } -func (r *hookRecorder) hook(_ context.Context, info InvocationInfo, failure *Failure) PostAction { +func (r *hookRecorder) hook(_ context.Context, info InvocationInfo, failure *Failure) FailureAction { r.mu.Lock() r.calls = append(r.calls, hookCall{info, failure}) r.mu.Unlock() @@ -81,11 +81,11 @@ func expectEventStream(mock *mockGrpcClient, events ...*pb.DispatchEvent) { }) } -func runWithPostHook(t *testing.T, handler any, timeoutMs int, rec *hookRecorder, setup ...func(*mockGrpcClient)) { +func runWithFailureHook(t *testing.T, handler any, timeoutMs int, rec *hookRecorder, setup ...func(*mockGrpcClient)) { controller := gomock.NewController(t) defer controller.Finish() executeHandlerWithAgent(t, controller, handler, timeoutMs, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.postDispatch = rec.hook + agent.onFailure = rec.hook expectEventStream(mock) for _, s := range setup { s(mock) @@ -93,41 +93,42 @@ func runWithPostHook(t *testing.T, handler any, timeoutMs int, rec *hookRecorder }) } -func TestPostDispatchHookOnSuccess(t *testing.T) { - rec := newHookRecorder(PostActionRetry) - runWithPostHook(t, func(ctx HandlerContext) error { return nil }, 0, rec) +func TestFailureHookNotCalledOnSuccess(t *testing.T) { + rec := newHookRecorder(FailureActionRetry) + runWithFailureHook(t, func(ctx HandlerContext) error { return nil }, 0, rec) + + rec.mu.Lock() + defer rec.mu.Unlock() + require.Empty(t, rec.calls) +} + +func TestFailureHookOnHandlerError(t *testing.T) { + rec := newHookRecorder(FailureActionIgnore) + runWithFailureHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec) call := rec.only(t) - require.Nil(t, call.failure) require.Equal(t, "func1", call.info.HandlerName) require.Equal(t, pb.HandlerInvokeType_RUN_NOW, call.info.InvokeType) require.Equal(t, 1, call.info.Attempt, "an agent that doesn't send attempt means the first") - // Retry on success is ignored: no RetryInvocation expectation is set. -} - -func TestPostDispatchHookOnHandlerError(t *testing.T) { - rec := newHookRecorder(PostActionIgnore) - runWithPostHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec) - - failure := rec.only(t).failure + failure := call.failure require.NotNil(t, failure) require.Equal(t, pb.FailureReason_HANDLER_ERROR, failure.Reason) require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) require.Equal(t, "ohnoes", failure.Message) } -func TestPostDispatchHookOnPanic(t *testing.T) { - rec := newHookRecorder(PostActionIgnore) - runWithPostHook(t, func(ctx HandlerContext) error { panic("ohnoes") }, 0, rec) +func TestFailureHookOnPanic(t *testing.T) { + rec := newHookRecorder(FailureActionIgnore) + runWithFailureHook(t, func(ctx HandlerContext) error { panic("ohnoes") }, 0, rec) failure := rec.only(t).failure require.Equal(t, pb.FailureReason_HANDLER_ERROR, failure.Reason) require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) } -func TestPostDispatchHookOnHandlerTimeout(t *testing.T) { - rec := newHookRecorder(PostActionIgnore) - runWithPostHook(t, func(ctx HandlerContext) error { +func TestFailureHookOnHandlerTimeout(t *testing.T) { + rec := newHookRecorder(FailureActionIgnore) + runWithFailureHook(t, func(ctx HandlerContext) error { time.Sleep(50 * time.Millisecond) return nil }, 1, rec) @@ -137,9 +138,9 @@ func TestPostDispatchHookOnHandlerTimeout(t *testing.T) { require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) } -func TestPostDispatchHookTimeoutDuringApiCallIsRuntime(t *testing.T) { - rec := newHookRecorder(PostActionIgnore) - runWithPostHook(t, func(ctx HandlerContext) error { +func TestFailureHookTimeoutDuringApiCallIsRuntime(t *testing.T) { + rec := newHookRecorder(FailureActionIgnore) + runWithFailureHook(t, func(ctx HandlerContext) error { _, err := ctx.CortexJsonApiCall("GET", "/api/v1/slow", "") return err }, 20, rec, func(mock *mockGrpcClient) { @@ -155,10 +156,10 @@ func TestPostDispatchHookTimeoutDuringApiCallIsRuntime(t *testing.T) { require.Equal(t, pb.FailureSource_RUNTIME, failure.Source) } -func TestPostDispatchHookRetry(t *testing.T) { - rec := newHookRecorder(PostActionRetry) +func TestFailureHookRetry(t *testing.T) { + rec := newHookRecorder(FailureActionRetry) retried := make(chan *pb.RetryInvocationRequest, 1) - runWithPostHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec, func(mock *mockGrpcClient) { + runWithFailureHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec, func(mock *mockGrpcClient) { mock.agentStub.EXPECT().RetryInvocation(gomock.Any(), gomock.Any()).DoAndReturn( func(ctx context.Context, req *pb.RetryInvocationRequest, opts ...grpc.CallOption) (*pb.RetryInvocationResponse, error) { retried <- req @@ -170,7 +171,7 @@ func TestPostDispatchHookRetry(t *testing.T) { require.Equal(t, "func1", req.Invoke.HandlerName) } -func TestPostDispatchHookPanicIsContained(t *testing.T) { +func TestFailureHookPanicIsContained(t *testing.T) { controller := gomock.NewController(t) defer controller.Finish() @@ -178,31 +179,14 @@ func TestPostDispatchHookPanicIsContained(t *testing.T) { executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, func(*pb.ReportInvocationRequest) { reported = true }, func(agent *Agent, mock *mockGrpcClient) { - agent.postDispatch = func(context.Context, InvocationInfo, *Failure) PostAction { panic("hook bug") } + agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { panic("hook bug") } expectEventStream(mock) }) require.True(t, reported) } -func TestPreDispatchHookRunsBeforeHandler(t *testing.T) { - controller := gomock.NewController(t) - defer controller.Finish() - - var order []string - executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { - order = append(order, "handler") - return nil - }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.preDispatch = func(_ context.Context, info InvocationInfo) { - require.Equal(t, "func1", info.HandlerName) - order = append(order, "pre") - } - }) - require.Equal(t, []string{"pre", "handler"}, order) -} - -func TestFailureEventReachesPostDispatchHook(t *testing.T) { - rec := newHookRecorder(PostActionRetry) +func TestFailureEventReachesFailureHook(t *testing.T) { + rec := newHookRecorder(FailureActionRetry) retried := make(chan *pb.RetryInvocationRequest, 1) ev := &pb.DispatchEvent{ @@ -228,7 +212,7 @@ func TestFailureEventReachesPostDispatchHook(t *testing.T) { <-rec.seen return nil }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.postDispatch = rec.hook + agent.onFailure = rec.hook expectEventStream(mock, ev) mock.agentStub.EXPECT().RetryInvocation(gomock.Any(), gomock.Any()).DoAndReturn( func(ctx context.Context, req *pb.RetryInvocationRequest, opts ...grpc.CallOption) (*pb.RetryInvocationResponse, error) { @@ -256,7 +240,7 @@ func TestFailureEventsStopOnOlderAgent(t *testing.T) { time.Sleep(50 * time.Millisecond) return nil }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.postDispatch = func(context.Context, InvocationInfo, *Failure) PostAction { return PostActionIgnore } + agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { return FailureActionIgnore } // Times(1): an agent without the RPC must not be redialed in a loop. mock.agentStub.EXPECT().DispatchEvents(gomock.Any(), gomock.Any()).Times(1).Return( &eventStream{err: status.Error(codes.Unimplemented, "unknown method DispatchEvents")}, nil) @@ -273,7 +257,7 @@ func TestDispatchRequestAdvertisesFailureEvents(t *testing.T) { executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { return nil }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { if withHook { - agent.postDispatch = func(context.Context, InvocationInfo, *Failure) PostAction { return PostActionIgnore } + agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { return FailureActionIgnore } expectEventStream(mock) } sent = mock.lastStream diff --git a/sdks/go/options.go b/sdks/go/options.go index e55bb148..3f60e18b 100644 --- a/sdks/go/options.go +++ b/sdks/go/options.go @@ -28,8 +28,7 @@ type agentOptions struct { sleepOnError time.Duration version string maxReceiveMessageSize int - preDispatch PreDispatchHook - postDispatch PostDispatchHook + onFailure FailureHook } func resolveMaxReceiveMessageSize() int { diff --git a/sdks/python/README.md b/sdks/python/README.md index fa6e20ee..a332d7b2 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -109,16 +109,15 @@ Set it to `1` for the old behaviour, where handlers run one at a time. ## Failure hooks -Pass `post_dispatch` to `AxonClient` to run your own code when an +Pass `on_failure` to `AxonClient` to run your own code whenever an invocation fails, for example to send the failed message somewhere you can -replay it from. The hook runs after every invocation (`failure` is `None` on -success). It also runs for failures that only the agent can see, such as a -webhook the agent couldn't hand to this process because every handler was -busy. Those failures arrive on their own stream, so they get through even +replay it from. It also runs for failures that only the agent can see, such +as a webhook the agent couldn't hand to this process because every handler +was busy. Those failures arrive on their own stream, so they get through even while every handler slot is taken. ```python -from cortex_axon import AxonClient, FailureSource, PostAction +from cortex_axon import AxonClient, FailureAction, FailureSource from cortex_axon.handler import cortex_webhook @@ -127,9 +126,7 @@ def my_webhook_handler(context): process(context.args["body"]) -def on_dispatch(info, failure): - if failure is None: - return # it worked +def on_failure(info, failure): # info.args holds the original webhook: "body", "content-type", "url" dead_letters.publish({ "handler": info.handler_name, @@ -142,30 +139,27 @@ def on_dispatch(info, failure): }) # Optionally have the agent run it again, as a new attempt. if failure.source == FailureSource.RUNTIME and info.attempt < 3: - return PostAction.RETRY + return FailureAction.RETRY -client = AxonClient(scope=globals(), post_dispatch=on_dispatch) +client = AxonClient(scope=globals(), on_failure=on_failure) client.run() ``` | Reason | When | Source | |---|---|---| | `HANDLER_ERROR` | The handler raised | `APPLICATION` | -| `TIMEOUT` | The handler ran past its timeout | `RUNTIME` if a Cortex API call was in flight or the agent never got a result; otherwise `APPLICATION` | +| `TIMEOUT` | The handler ran past its timeout | `RUNTIME` if a Cortex API call was in flight, the invocation waited for a free handler, or the agent never got a result; otherwise `APPLICATION` | | `NO_RESOURCES` | The agent's queue was full, or the message expired waiting for a free handler | `RUNTIME` | `caller_rejected` is true when the sender was told the message wasn't accepted (a webhook got a 503), so the sender may also retry it. -Returning `PostAction.RETRY` asks the agent to queue the invocation again, -with `attempt` incremented. The agent stops after 5 attempts. Returning -nothing ignores the failure. +Returning `FailureAction.RETRY` asks the agent to queue the invocation +again, with `attempt` incremented. The agent stops after 5 attempts. +Returning nothing ignores the failure. -`pre_dispatch` takes a function of `info` that runs before each handler, -e.g. to record every inbound message before processing it. - -Hooks run for every handler type, not only webhooks. Python can't stop a +The hook runs for every handler type, not only webhooks. Python can't stop a handler that overruns its timeout: at the deadline the SDK reports the `TIMEOUT` and runs the hook, and whatever the handler does after that is dropped. diff --git a/sdks/python/cortex_axon/__init__.py b/sdks/python/cortex_axon/__init__.py index 1150881e..f31cadd4 100644 --- a/sdks/python/cortex_axon/__init__.py +++ b/sdks/python/cortex_axon/__init__.py @@ -1,7 +1,7 @@ from .axon_client import AxonClient as AxonClient from .axon_client import CortexHTTPError as CortexHTTPError from .hooks import Failure as Failure +from .hooks import FailureAction as FailureAction from .hooks import FailureReason as FailureReason from .hooks import FailureSource as FailureSource from .hooks import InvocationInfo as InvocationInfo -from .hooks import PostAction as PostAction diff --git a/sdks/python/cortex_axon/axon_client.py b/sdks/python/cortex_axon/axon_client.py index e0c76735..536a18cc 100644 --- a/sdks/python/cortex_axon/axon_client.py +++ b/sdks/python/cortex_axon/axon_client.py @@ -26,12 +26,11 @@ ) from .hooks import ( Failure, + FailureAction, + FailureHook, FailureReason, FailureSource, InvocationInfo, - PostAction, - PostDispatchHook, - PreDispatchHook, ) DEFAULT_MAX_CONCURRENT_HANDLERS = 8 @@ -181,8 +180,7 @@ def __init__( scope=None, max_receive_message_size: Optional[int] = None, max_concurrent_handlers: int = DEFAULT_MAX_CONCURRENT_HANDLERS, - pre_dispatch: Optional[PreDispatchHook] = None, - post_dispatch: Optional[PostDispatchHook] = None, + on_failure: Optional[FailureHook] = None, ): if max_concurrent_handlers < 1: raise ValueError("max_concurrent_handlers must be at least 1") @@ -227,8 +225,7 @@ def __init__( ) self._handler_slots = threading.BoundedSemaphore(max_concurrent_handlers) - self.pre_dispatch = pre_dispatch - self.post_dispatch = post_dispatch + self.on_failure = on_failure # Failure events have their own thread and pool, apart from the # handler slots: they arrive exactly when every handler is busy, and # must never wait for one to free up. @@ -236,7 +233,7 @@ def __init__( self._events_thread = None self._events_stop = threading.Event() self._events_call = None - if post_dispatch: + if on_failure: self._hook_executor = ThreadPoolExecutor( max_workers=2, thread_name_prefix="axon-hook" ) @@ -276,25 +273,16 @@ def _find_handler(self, name): return h return None - def _run_pre_dispatch(self, info: InvocationInfo): - if not self.pre_dispatch: + def _run_failure_hook(self, info: InvocationInfo, failure: Failure, stub): + if not self.on_failure: return try: - self.pre_dispatch(info) + action = self.on_failure(info, failure) except Exception: - print(f"Error in pre_dispatch hook for {info.handler_name} (id={info.invocation_id})") - traceback.print_exc() - - def _run_post_dispatch(self, info: InvocationInfo, failure: Optional[Failure], stub): - if not self.post_dispatch: - return - try: - action = self.post_dispatch(info, failure) - except Exception: - print(f"Error in post_dispatch hook for {info.handler_name} (id={info.invocation_id})") + print(f"Error in on_failure hook for {info.handler_name} (id={info.invocation_id})") traceback.print_exc() return - if action == PostAction.RETRY and failure is not None: + if action == FailureAction.RETRY: self._retry_invocation(info, stub) def _retry_invocation(self, info: InvocationInfo, stub): @@ -317,7 +305,7 @@ def _retry_invocation(self, info: InvocationInfo, stub): ) def _start_failure_events(self): - if not self.post_dispatch or self._events_thread: + if not self.on_failure or self._events_thread: return self._events_thread = threading.Thread( target=self._run_failure_events, name="axon-failure-events", daemon=True @@ -351,7 +339,7 @@ def _run_failure_events(self): if not event.HasField("invoke"): continue self._hook_executor.submit( - self._run_post_dispatch, + self._run_failure_hook, InvocationInfo.from_proto(event.invoke), Failure.from_proto(event.failure), stub, @@ -404,7 +392,6 @@ def _invoke_handler(self, handler_to_invoke, invoke, stub): failure = None result = None info = InvocationInfo.from_proto(invoke) - self._run_pre_dispatch(info) now = datetime.now() timeout, capped, _ = _invoke_timeout(invoke, time()) ctx = self._handler_context(invoke.args, timeout * 1000 if timeout else 0) @@ -443,7 +430,7 @@ def on_timeout(): error=common_pb2.Error(code="timeout", message=message), logs=list(ctx.logs), )) - self._run_post_dispatch( + self._run_failure_hook( info, Failure(reason=FailureReason.TIMEOUT, source=source, message=message), stub ) @@ -501,7 +488,8 @@ def on_timeout(): # After the report, so a slow hook can't make the agent decide this # invocation was never answered and report the failure a second time. - self._run_post_dispatch(info, failure, stub) + if failure is not None: + self._run_failure_hook(info, failure, stub) def _report(self, stub, invoke_info): # A failure here must not escape the worker: the dispatch loop is no @@ -538,7 +526,7 @@ def _run_dispatch(self): initial_message = cortex_axon_agent_pb2.DispatchRequest( dispatch_id=self.id, client_version=AGENT_VERSION, - failure_events=self.post_dispatch is not None, + failure_events=self.on_failure is not None, ) for response in stub.Dispatch( _client_message_iterator(initial_message) diff --git a/sdks/python/cortex_axon/hooks.py b/sdks/python/cortex_axon/hooks.py index 6b910c54..f4240d8c 100644 --- a/sdks/python/cortex_axon/hooks.py +++ b/sdks/python/cortex_axon/hooks.py @@ -24,7 +24,7 @@ class FailureSource(IntEnum): RUNTIME = cortex_axon_agent_pb2.RUNTIME -class PostAction(Enum): +class FailureAction(Enum): # Leave the failure as it is. IGNORE = "ignore" # Queue the invocation again as its next attempt. @@ -75,11 +75,7 @@ def from_proto(cls, failure: cortex_axon_agent_pb2.InvocationFailure) -> "Failur ) -# Runs before each handler. -PreDispatchHook = Callable[[InvocationInfo], None] - -# Runs after each invocation with failure=None on success, and for every -# failure the agent detects on its own, such as an invocation that never -# reached this process. Returning None means PostAction.IGNORE; the return -# value only matters for a failure. -PostDispatchHook = Callable[[InvocationInfo, Optional[Failure]], Optional[PostAction]] +# Runs for every failed invocation, including failures the agent detects on +# its own, such as an invocation that never reached this process. Returning +# None means FailureAction.IGNORE. +FailureHook = Callable[[InvocationInfo, Failure], Optional[FailureAction]] diff --git a/sdks/python/example/main.py b/sdks/python/example/main.py index ca92ffe4..f4ded6a4 100644 --- a/sdks/python/example/main.py +++ b/sdks/python/example/main.py @@ -1,6 +1,6 @@ import json -from cortex_axon import PostAction +from cortex_axon import FailureAction from cortex_axon.axon_client import AxonClient, HandlerContext from cortex_axon.handler import cortex_handler, cortex_scheduled @@ -51,25 +51,23 @@ def my_invoke_handler(context: HandlerContext) -> str: return json.dumps(result) -def on_dispatch(info, failure): - # Runs after every invocation, and for failures only the agent sees. - if failure is None: - return None +def on_failure(info, failure): + # Runs for every failed invocation, including ones only the agent sees. print( f"FAILED {info.handler_name} (id={info.invocation_id}, attempt={info.attempt}) " f"reason={failure.reason.name} source={failure.source.name}: {failure.message}" ) # A real hook would save info.args somewhere it can be replayed from. if info.attempt < 3: - return PostAction.RETRY - return PostAction.IGNORE + return FailureAction.RETRY + return FailureAction.IGNORE def run(): # Connect to the gRPC server client = AxonClient( scope=globals(), - post_dispatch=on_dispatch, + on_failure=on_failure, ) client.run() diff --git a/sdks/python/tests/test_hooks.py b/sdks/python/tests/test_hooks.py index 360e7953..1f23b524 100644 --- a/sdks/python/tests/test_hooks.py +++ b/sdks/python/tests/test_hooks.py @@ -9,9 +9,9 @@ from cortex_axon import ( AxonClient, + FailureAction, FailureReason, FailureSource, - PostAction, ) from cortex_axon.handler import CortexWebhook from generated import ( @@ -106,14 +106,10 @@ def serve(servicer, add): class Hooks: def __init__(self, action=None): self.action = action - self.pre = queue.Queue() - self.post = queue.Queue() + self.failures = queue.Queue() - def pre_dispatch(self, info): - self.pre.put(info) - - def post_dispatch(self, info, failure): - self.post.put((info, failure)) + def on_failure(self, info, failure): + self.failures.put((info, failure)) return self.action @@ -141,8 +137,7 @@ def start(agent, handler, hooks=None, api=None, max_concurrent_handlers=4): cortex_port=api_port, handlers=[webhook], max_concurrent_handlers=max_concurrent_handlers, - pre_dispatch=hooks.pre_dispatch if hooks else None, - post_dispatch=hooks.post_dispatch if hooks else None, + on_failure=hooks.on_failure if hooks else None, ) threading.Thread(target=client.run, daemon=True).start() clients.append(client) @@ -164,8 +159,8 @@ def start(agent, handler, hooks=None, api=None, max_concurrent_handlers=4): client._events_thread.join(timeout=5) -def test_post_dispatch_on_success(run_client): - hooks = Hooks(action=PostAction.RETRY) +def test_failure_hook_not_called_on_success(run_client): + hooks = Hooks(action=FailureAction.RETRY) agent = HookAgent(invokes=[invoke(body="payload")]) def my_hooked_webhook(ctx): @@ -173,19 +168,14 @@ def my_hooked_webhook(ctx): run_client(agent, my_hooked_webhook, hooks) - pre = hooks.pre.get(timeout=5) - info, failure = hooks.post.get(timeout=5) - assert pre.invocation_id == info.invocation_id == "1" - assert failure is None - assert info.args == {"body": "payload"} - assert info.attempt == 1 - # RETRY on success is ignored. + assert agent.reports.get(timeout=5).result.value == "ok" with pytest.raises(queue.Empty): - agent.retries.get(timeout=0.5) + hooks.failures.get(timeout=0.5) + assert agent.retries.empty() -def test_post_dispatch_on_handler_error_and_retry(run_client): - hooks = Hooks(action=PostAction.RETRY) +def test_failure_hook_on_handler_error_and_retry(run_client): + hooks = Hooks(action=FailureAction.RETRY) agent = HookAgent(invokes=[invoke(body="payload")]) def my_hooked_webhook(ctx): @@ -193,7 +183,10 @@ def my_hooked_webhook(ctx): run_client(agent, my_hooked_webhook, hooks) - info, failure = hooks.post.get(timeout=5) + info, failure = hooks.failures.get(timeout=5) + assert info.invocation_id == "1" + assert info.args == {"body": "payload"} + assert info.attempt == 1 assert failure.reason == FailureReason.HANDLER_ERROR assert failure.source == FailureSource.APPLICATION assert failure.message == "ohnoes" @@ -213,7 +206,7 @@ def my_hooked_webhook(ctx): try: run_client(agent, my_hooked_webhook, hooks, api=api) - _, failure = hooks.post.get(timeout=5) + _, failure = hooks.failures.get(timeout=5) finally: api.release.set() assert failure.reason == FailureReason.TIMEOUT @@ -231,7 +224,7 @@ def my_hooked_webhook(ctx): run_client(agent, my_hooked_webhook, hooks) - _, failure = hooks.post.get(timeout=5) + _, failure = hooks.failures.get(timeout=5) assert failure.reason == FailureReason.TIMEOUT assert failure.source == FailureSource.APPLICATION report = agent.reports.get(timeout=5) @@ -240,7 +233,7 @@ def my_hooked_webhook(ctx): # The late finish is dropped: no second report, no second hook call. assert finished.wait(timeout=5) with pytest.raises(queue.Empty): - hooks.post.get(timeout=0.5) + hooks.failures.get(timeout=0.5) assert agent.reports.empty() @@ -259,12 +252,11 @@ def my_hooked_webhook(ctx): run_client(agent, my_hooked_webhook, hooks) - info, failure = hooks.post.get(timeout=5) - assert info.invocation_id == "fresh" - assert failure is None assert agent.reports.get(timeout=5).handler_invoke.invocation_id == "fresh" + # The stale one neither ran nor reached the hook: the agent reports it. with pytest.raises(queue.Empty): - hooks.post.get(timeout=0.5) + hooks.failures.get(timeout=0.5) + assert agent.reports.empty() assert ran.qsize() == 1 @@ -278,7 +270,7 @@ def my_hooked_webhook(ctx): run_client(agent, my_hooked_webhook, hooks) - _, failure = hooks.post.get(timeout=5) + _, failure = hooks.failures.get(timeout=5) assert failure.reason == FailureReason.TIMEOUT assert failure.source == FailureSource.RUNTIME assert agent.reports.get(timeout=5).error.code == "timeout" @@ -289,15 +281,19 @@ def test_hook_errors_are_contained(run_client): ran = queue.Queue() class Broken(Hooks): - def post_dispatch(self, info, failure): + def on_failure(self, info, failure): raise RuntimeError("hook bug") def my_hooked_webhook(ctx): ran.put(True) + raise ValueError("ohnoes") run_client(agent, my_hooked_webhook, Broken()) + # Both still run and report, despite the hook raising each time. ran.get(timeout=5) ran.get(timeout=5) + assert agent.reports.get(timeout=5).error.code == "unexpected" + assert agent.reports.get(timeout=5).error.code == "unexpected" def test_dispatch_request_advertises_failure_events(run_client): @@ -331,7 +327,7 @@ def test_failure_events_reach_hook_while_handlers_are_saturated(run_client): ] # More invocations than slots, so the dispatch reader blocks. agent = HookAgent(invokes=[invoke(str(i)) for i in range(slots + 2)], events=events) - hooks = Hooks(action=PostAction.RETRY) + hooks = Hooks(action=FailureAction.RETRY) def my_hooked_webhook(ctx): with contextlib.suppress(threading.BrokenBarrierError): @@ -342,7 +338,7 @@ def my_hooked_webhook(ctx): run_client(agent, my_hooked_webhook, hooks, max_concurrent_handlers=slots) started.wait(timeout=5) - seen = [hooks.post.get(timeout=5) for _ in events] + seen = [hooks.failures.get(timeout=5) for _ in events] bodies = sorted(info.args["body"] for info, _ in seen) assert bodies == sorted(f"msg-{i}" for i in range(20)) for _, failure in seen: From 5103af67ebe12be5e8ae1c4e8802456dd033479f Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Wed, 11 Nov 2026 11:24:01 +1100 Subject: [PATCH 4/5] End-to-end failure tests: SDKs against a real agent Both SDKs gain tests that build the agent, run it, and trigger invocations through its webhook endpoint: a handler error and a timeout each reach the failure hook and succeed on retry, and in Python, invocations stuck behind a saturated handler are reported by the agent, retried, and each succeed exactly once without the expired originals running. To support them: the agent's webhook port can be set with WEBHOOK_PORT, and the Python SDK workflow sets up Go to build the agent. The Go SDK's Run(ctx) now stops when ctx is cancelled; it had opened its dispatch stream with context.Background(), and its exit flag was raced. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/sdks-python-test.yml | 9 + agent/config/config.go | 11 +- sdks/go/axon_agent.go | 15 +- sdks/go/e2e_failure_test.go | 253 +++++++++++++++++++++++++ sdks/python/tests/test_e2e_failures.py | 243 ++++++++++++++++++++++++ 5 files changed, 523 insertions(+), 8 deletions(-) create mode 100644 sdks/go/e2e_failure_test.go create mode 100644 sdks/python/tests/test_e2e_failures.py diff --git a/.github/workflows/sdks-python-test.yml b/.github/workflows/sdks-python-test.yml index e1c233dd..168d4deb 100644 --- a/.github/workflows/sdks-python-test.yml +++ b/.github/workflows/sdks-python-test.yml @@ -34,6 +34,15 @@ jobs: sudo apt-get update sudo apt-get install -y protobuf-compiler + # The end-to-end tests build and run the real agent. + - name: Set up Go + uses: actions/setup-go@v3 + with: + go-version: '1.26.10' + + - name: Generate agent protobuf + run: make -C agent setup proto + - name: Run Linter run: | cd sdks/python diff --git a/agent/config/config.go b/agent/config/config.go index 7e0e71be..f2c2e76a 100644 --- a/agent/config/config.go +++ b/agent/config/config.go @@ -258,6 +258,15 @@ func NewAgentEnvConfig() AgentConfig { httpPort = p } + webhookPort := WebhookServerPort + if webhookPortStr := os.Getenv("WEBHOOK_PORT"); webhookPortStr != "" { + p, err := strconv.Atoi(webhookPortStr) + if err != nil { + panic(err) + } + webhookPort = p + } + snykBrokerPort := 0 if snykBrokerPortStr := os.Getenv("SNYK_BROKER_PORT"); snykBrokerPortStr != "" { p, err := strconv.Atoi(snykBrokerPortStr) @@ -337,7 +346,7 @@ func NewAgentEnvConfig() AgentConfig { InstanceId: getInstanceId(), IntegrationAlias: identifier, HttpServerPort: httpPort, - WebhookServerPort: WebhookServerPort, + WebhookServerPort: webhookPort, SnykBrokerPort: snykBrokerPort, EnableApiProxy: true, EnablePprof: enablePprof, diff --git a/sdks/go/axon_agent.go b/sdks/go/axon_agent.go index 7c0cbd96..c1c195dd 100644 --- a/sdks/go/axon_agent.go +++ b/sdks/go/axon_agent.go @@ -9,6 +9,7 @@ import ( "runtime" "strings" "sync" + "sync/atomic" "time" pb "github.com/cortexapps/axon-go/.generated/proto/github.com/cortexapps/axon" @@ -234,12 +235,12 @@ const sleepErrorWait = 5 * time.Second // Run starts the agent and begins dispatching invocations to the registered handlers. func (a *Agent) Run(ctx context.Context) error { - exit := false + var exit atomic.Bool reregister := false sleepOnError := func(err error) { if a.sleepOnError == 0 { - exit = true + exit.Store(true) return } @@ -270,13 +271,13 @@ func (a *Agent) Run(ctx context.Context) error { go func() { select { case <-ctx.Done(): - exit = true + exit.Store(true) case <-a.done: - exit = true + exit.Store(true) } }() - for !exit { + for !exit.Load() { // aquire an agent and register handlers if needed // this allows agent crash/restart to recover @@ -297,7 +298,7 @@ func (a *Agent) Run(ctx context.Context) error { reregister = false - stream, err := stub.Dispatch(context.Background()) + stream, err := stub.Dispatch(ctx) if err != nil { sleepOnError(err) continue @@ -371,7 +372,7 @@ func (a *Agent) processDispatchStream(ctx context.Context, stream grpc.BidiStrea req, err := stream.Recv() if err != nil { status, _ := status.FromError(err) - if err == io.EOF || status.Code() == codes.Unavailable { + if err == io.EOF || status.Code() == codes.Unavailable || ctx.Err() != nil { return nil } diff --git a/sdks/go/e2e_failure_test.go b/sdks/go/e2e_failure_test.go new file mode 100644 index 00000000..969086ee --- /dev/null +++ b/sdks/go/e2e_failure_test.go @@ -0,0 +1,253 @@ +package axon + +import ( + "context" + "fmt" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + pb "github.com/cortexapps/axon-go/.generated/proto/github.com/cortexapps/axon" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +// These tests run the SDK against a real agent built from ../../agent, and +// trigger invocations through its webhook endpoint. + +var ( + agentBinOnce sync.Once + agentBin string + agentBinErr error +) + +func buildAgent(t *testing.T) string { + agentBinOnce.Do(func() { + dir, err := os.MkdirTemp("", "axon-e2e") + if err != nil { + agentBinErr = err + return + } + agentBin = filepath.Join(dir, "axon-agent") + cmd := exec.Command("go", "build", "-o", agentBin, ".") + cmd.Dir = filepath.Join("..", "..", "agent") + out, err := cmd.CombinedOutput() + if err != nil { + agentBinErr = fmt.Errorf("building agent: %v\n%s", err, out) + } + }) + require.NoError(t, agentBinErr) + return agentBin +} + +func freePort(t *testing.T) int { + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + return ln.Addr().(*net.TCPAddr).Port +} + +type realAgent struct { + grpcPort int + webhookPort int +} + +func startRealAgent(t *testing.T) *realAgent { + if testing.Short() { + t.Skip("starts a real agent") + } + bin := buildAgent(t) + a := &realAgent{grpcPort: freePort(t), webhookPort: freePort(t)} + + cmd := exec.Command(bin, "serve") + cmd.Env = append(os.Environ(), + "DRYRUN=true", + fmt.Sprintf("PORT=%d", a.grpcPort), + fmt.Sprintf("HTTP_PORT=%d", freePort(t)), + fmt.Sprintf("WEBHOOK_PORT=%d", a.webhookPort), + "HANDLER_HISTORY_PATH="+t.TempDir(), + ) + logFile, err := os.Create(filepath.Join(t.TempDir(), "agent.log")) + require.NoError(t, err) + cmd.Stdout, cmd.Stderr = logFile, logFile + require.NoError(t, cmd.Start()) + t.Cleanup(func() { + _ = cmd.Process.Kill() + _ = cmd.Wait() + if t.Failed() { + logs, _ := os.ReadFile(logFile.Name()) + t.Logf("agent log:\n%s", logs) + } + }) + + require.Eventually(t, func() bool { + conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", a.grpcPort)) + if err == nil { + conn.Close() + } + return err == nil + }, 30*time.Second, 50*time.Millisecond, "agent never started listening") + return a +} + +// postWebhook sends body once the webhook is live, and returns the status the +// agent accepted it with. +func (a *realAgent) postWebhook(t *testing.T, id, body string) int { + url := fmt.Sprintf("http://127.0.0.1:%d/webhook/%s", a.webhookPort, id) + var status int + require.Eventually(t, func() bool { + resp, err := http.Post(url, "application/json", strings.NewReader(body)) + if err != nil { + return false + } + resp.Body.Close() + status = resp.StatusCode + // 404 until the SDK has registered and started dispatching. + return status != http.StatusNotFound + }, 10*time.Second, 50*time.Millisecond, "webhook never became available") + return status +} + +type e2eCall struct { + info InvocationInfo + failure *Failure +} + +// runAgainstRealAgent registers handler as webhook id on a fresh agent, with +// a failure hook that records each failure and asks for a retry. +func runAgainstRealAgent(t *testing.T, id string, handler Handler, opts ...RegisterHandlerOption) (*realAgent, chan e2eCall) { + a := startRealAgent(t) + failures := make(chan e2eCall, 10) + + logConfig := zap.NewDevelopmentConfig() + logConfig.Level = zap.NewAtomicLevelAt(zap.WarnLevel) + agent := NewAxonAgent( + WithHostport("127.0.0.1", a.grpcPort), + WithLoggerConfig(logConfig), + WithFailureHook(func(_ context.Context, info InvocationInfo, failure *Failure) FailureAction { + failures <- e2eCall{info, failure} + return FailureActionRetry + }), + ) + + opts = append(opts, WithInvokeOption(pb.HandlerInvokeType_WEBHOOK, id)) + _, err := agent.RegisterHandler(handler, opts...) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + _ = agent.Run(ctx) + }() + t.Cleanup(func() { + cancel() + select { + case <-done: + case <-time.After(10 * time.Second): + } + }) + return a, failures +} + +type attempts struct { + mu sync.Mutex + bodies []string + ok chan string +} + +func newAttempts() *attempts { return &attempts{ok: make(chan string, 10)} } + +func (r *attempts) record(ctx HandlerContext) int { + r.mu.Lock() + defer r.mu.Unlock() + r.bodies = append(r.bodies, ctx.Args()["body"]) + return len(r.bodies) +} + +func requireFailure(t *testing.T, failures chan e2eCall) e2eCall { + select { + case call := <-failures: + return call + case <-time.After(15 * time.Second): + require.FailNow(t, "failure hook never ran") + return e2eCall{} + } +} + +func requireSucceeded(t *testing.T, r *attempts) string { + select { + case body := <-r.ok: + return body + case <-time.After(15 * time.Second): + require.FailNow(t, "retry never succeeded") + return "" + } +} + +func TestE2EHandlerErrorIsHookedAndRetried(t *testing.T) { + r := newAttempts() + a, failures := runAgainstRealAgent(t, "e2e-error", func(ctx HandlerContext) error { + if r.record(ctx) == 1 { + return fmt.Errorf("first attempt fails") + } + r.ok <- ctx.Args()["body"] + return nil + }) + + require.Equal(t, http.StatusOK, a.postWebhook(t, "e2e-error", `{"id": 42}`)) + + call := requireFailure(t, failures) + require.Equal(t, pb.FailureReason_HANDLER_ERROR, call.failure.Reason) + require.Equal(t, pb.FailureSource_APPLICATION, call.failure.Source) + require.Equal(t, "first attempt fails", call.failure.Message) + require.Equal(t, `{"id": 42}`, call.info.Args["body"]) + require.Equal(t, 1, call.info.Attempt) + + require.Equal(t, `{"id": 42}`, requireSucceeded(t, r), "the retry carries the original message") + + // It succeeded, so the hook does not run again. + select { + case extra := <-failures: + require.FailNow(t, "unexpected second failure", "%+v", extra.failure) + case <-time.After(time.Second): + } + r.mu.Lock() + require.Len(t, r.bodies, 2) + r.mu.Unlock() +} + +func TestE2ETimeoutIsHookedAndRetried(t *testing.T) { + r := newAttempts() + a, failures := runAgainstRealAgent(t, "e2e-timeout", func(ctx HandlerContext) error { + if r.record(ctx) == 1 { + <-ctx.Done() + return ctx.Err() + } + r.ok <- ctx.Args()["body"] + return nil + }, WithTimeout(500*time.Millisecond)) + + require.Equal(t, http.StatusOK, a.postWebhook(t, "e2e-timeout", `{"id": 7}`)) + + call := requireFailure(t, failures) + require.Equal(t, pb.FailureReason_TIMEOUT, call.failure.Reason) + require.Equal(t, pb.FailureSource_APPLICATION, call.failure.Source) + require.Equal(t, `{"id": 7}`, call.info.Args["body"]) + + require.Equal(t, `{"id": 7}`, requireSucceeded(t, r)) + + // The SDK reported the timeout itself, so the agent must not report it + // again once its grace period is over. + select { + case extra := <-failures: + require.FailNow(t, "timeout reported twice", "%+v", extra.failure) + case <-time.After(7 * time.Second): + } +} diff --git a/sdks/python/tests/test_e2e_failures.py b/sdks/python/tests/test_e2e_failures.py new file mode 100644 index 00000000..026a10e7 --- /dev/null +++ b/sdks/python/tests/test_e2e_failures.py @@ -0,0 +1,243 @@ +"""The SDK against a real agent, triggered through the agent's webhook +endpoint. The agent is built from ../../agent with `go`, or taken from +AXON_AGENT_BIN; without either these tests are skipped.""" + +import collections +import os +import queue +import shutil +import socket +import subprocess +import threading +import time +import urllib.error +import urllib.request +from pathlib import Path + +import pytest + +from cortex_axon import AxonClient, FailureAction, FailureReason, FailureSource +from cortex_axon.handler import CortexWebhook + +AGENT_DIR = Path(__file__).resolve().parents[3] / "agent" + +# AxonClient.run has no way to stop, so a test's client keeps retrying after +# its agent is killed, and may print into pytest's closed capture. +pytestmark = pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") + + +@pytest.fixture(scope="session") +def agent_bin(tmp_path_factory): + if os.environ.get("AXON_AGENT_BIN"): + return os.environ["AXON_AGENT_BIN"] + if not shutil.which("go"): + pytest.skip("needs go to build the agent, or AXON_AGENT_BIN") + out = tmp_path_factory.mktemp("agent") / "axon-agent" + subprocess.run(["go", "build", "-o", str(out), "."], cwd=AGENT_DIR, check=True) + return str(out) + + +def free_port(): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class RealAgent: + def __init__(self, agent_bin, tmp_path): + self.grpc_port = free_port() + self.webhook_port = free_port() + env = dict( + os.environ, + DRYRUN="true", + PORT=str(self.grpc_port), + HTTP_PORT=str(free_port()), + WEBHOOK_PORT=str(self.webhook_port), + HANDLER_HISTORY_PATH=str(tmp_path / "history"), + ) + # Stays open for the agent process's lifetime; closed in stop(). + self.log = open(tmp_path / "agent.log", "wb") # noqa: SIM115 + self.proc = subprocess.Popen( + [agent_bin, "serve"], env=env, stdout=self.log, stderr=subprocess.STDOUT + ) + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + try: + socket.create_connection(("127.0.0.1", self.grpc_port), timeout=0.2).close() + return + except OSError: + time.sleep(0.05) + raise AssertionError("agent never started listening") + + def post(self, webhook_id, body): + """Sends body once the webhook is live; returns the HTTP status.""" + url = f"http://127.0.0.1:{self.webhook_port}/webhook/{webhook_id}" + deadline = time.monotonic() + 10 + while True: + req = urllib.request.Request(url, data=body.encode(), method="POST", + headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req) as r: + return r.status + except urllib.error.HTTPError as e: + # 404 until the SDK has registered and started dispatching. + if e.code != 404 or time.monotonic() > deadline: + return e.code + time.sleep(0.05) + + def stop(self): + self.proc.kill() + self.proc.wait() + self.log.close() + + +@pytest.fixture +def real_agent(agent_bin, tmp_path): + agent = RealAgent(agent_bin, tmp_path) + yield agent + agent.stop() + + +class Recorder: + """Records each run of the handler and each call of the failure hook, + which always asks for a retry.""" + + def __init__(self): + self.lock = threading.Lock() + self.runs = collections.Counter() + self.succeeded = collections.Counter() + self.failures = [] + self.failure_seen = queue.Queue() + self.success_seen = queue.Queue() + + def ran(self, ctx): + with self.lock: + self.runs[ctx.args["body"]] += 1 + return self.runs[ctx.args["body"]] + + def ok(self, ctx): + with self.lock: + self.succeeded[ctx.args["body"]] += 1 + self.success_seen.put(ctx.args["body"]) + + def on_failure(self, info, failure): + with self.lock: + self.failures.append((info, failure)) + self.failure_seen.put((info, failure)) + return FailureAction.RETRY + + +def run_client(agent, recorder, webhook_id, func, timeout_ms=0, slots=4): + webhook = CortexWebhook(id=webhook_id, timeout_ms=timeout_ms) + webhook.set_func(func) + client = AxonClient( + agent_host="127.0.0.1", + agent_port=agent.grpc_port, + handlers=[webhook], + max_concurrent_handlers=slots, + on_failure=recorder.on_failure, + ) + threading.Thread(target=client.run, daemon=True).start() + return client + + +def test_handler_error_is_hooked_and_retried(real_agent): + rec = Recorder() + + def e2e_error(ctx): + if rec.ran(ctx) == 1: + raise ValueError("first attempt fails") + rec.ok(ctx) + + run_client(real_agent, rec, "e2e-error", e2e_error) + assert real_agent.post("e2e-error", '{"id": 42}') == 200 + + info, failure = rec.failure_seen.get(timeout=15) + assert failure.reason == FailureReason.HANDLER_ERROR + assert failure.source == FailureSource.APPLICATION + assert failure.message == "first attempt fails" + assert info.args["body"] == '{"id": 42}' + assert info.attempt == 1 + + # The retry carries the original message, and succeeds. + assert rec.success_seen.get(timeout=15) == '{"id": 42}' + with pytest.raises(queue.Empty): + rec.failure_seen.get(timeout=1) + assert rec.runs['{"id": 42}'] == 2 + + +def test_timeout_is_hooked_and_retried(real_agent): + rec = Recorder() + + def e2e_timeout(ctx): + if rec.ran(ctx) == 1: + time.sleep(3) + return + rec.ok(ctx) + + run_client(real_agent, rec, "e2e-timeout", e2e_timeout, timeout_ms=500) + assert real_agent.post("e2e-timeout", '{"id": 7}') == 200 + + info, failure = rec.failure_seen.get(timeout=15) + assert failure.reason == FailureReason.TIMEOUT + assert failure.source == FailureSource.APPLICATION + assert info.args["body"] == '{"id": 7}' + + assert rec.success_seen.get(timeout=15) == '{"id": 7}' + + # The SDK reported the timeout itself, so the agent must not report it + # again once its grace period is over, and the late finish is dropped. + with pytest.raises(queue.Empty): + rec.failure_seen.get(timeout=7) + assert rec.succeeded['{"id": 7}'] == 1 + + +def test_saturated_invocations_are_hooked_retried_and_run_once(real_agent): + """One handler slot, stuck. The invocations queued behind it expire; the + agent reports them, the hook retries them, and each message ends up + succeeding exactly once. None of the expired originals runs late.""" + rec = Recorder() + release = threading.Event() + bodies = [f"m{i}" for i in range(5)] + + def e2e_busy(ctx): + rec.ran(ctx) + if not release.is_set(): + release.wait(30) + return + rec.ok(ctx) + + original_hook = rec.on_failure + + def on_failure(info, failure): + if failure.source == FailureSource.RUNTIME: + # The agent has given up on the queued ones: unstick the slot. + release.set() + return original_hook(info, failure) + + rec.on_failure = on_failure + run_client(real_agent, rec, "e2e-busy", e2e_busy, timeout_ms=1000, slots=1) + + assert real_agent.post("e2e-busy", bodies[0]) == 200 + for body in bodies[1:]: + assert real_agent.post("e2e-busy", body) == 200 + + succeeded = {rec.success_seen.get(timeout=30) for _ in bodies} + assert succeeded == set(bodies) + with pytest.raises(queue.Empty): + rec.success_seen.get(timeout=2) + + with rec.lock: + assert all(rec.succeeded[b] == 1 for b in bodies), rec.succeeded + # Each failure is reported once per invocation. + ids = [info.invocation_id for info, _ in rec.failures] + assert len(ids) == len(set(ids)), "an invocation was reported twice" + # Every message failed at least once, and the queued ones were + # reported by the agent as Axon's fault. + failed = collections.Counter(info.args["body"] for info, _ in rec.failures) + assert set(failed) == set(bodies) + runtime = {info.args["body"] for info, f in rec.failures if f.source == FailureSource.RUNTIME} + assert set(bodies[1:]) <= runtime + # The expired originals never ran: each queued message ran once, as + # its retry. + assert all(rec.runs[b] == 1 for b in bodies[1:]), rec.runs From 23b4a532c51ee067c07447ea02cca1f6bc8c65a3 Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Wed, 11 Nov 2026 11:52:57 +1100 Subject: [PATCH 5/5] Route every failure through a FailureManager; the hook answers each one Addresses review feedback: - Invocations have one SetState with an explicit transition table, instead of ClaimDispatch/ClaimExpired. - Failure handling moves out of Manager into FailureManager, which owns the queues, pending answers, metric, retries and retry cap. Retries go back through Manager.Trigger. - One path into it: every failure completes its invocation with a FailureError, and one watcher in Manager hands those to FailureManager. That covers a full queue, expiry in the queue, the agent's own timeout, and failures the SDK reports, which now carry their classification on ReportInvocationRequest.failure. - The send loop is generic: drain() sends whatever an InvocationSource hands out. Dispatch drains Manager and DispatchEvents drains FailureManager. - SDKs no longer call RetryInvocation, which is removed. DispatchEvents is bidirectional: the agent sends each failure, the SDK runs the hook and answers DROP or RETRY on the same stream, and the agent applies it. An event unanswered when the stream ends is sent again on the next one. Co-Authored-By: Claude Opus 5.5 --- agent/cmd/relay.go | 1 + agent/proto/cortex-axon-agent.proto | 41 ++- agent/server/failure_events_e2e_test.go | 104 ++++-- agent/server/handler/failure_events_test.go | 212 +++++++---- agent/server/handler/failure_manager.go | 258 ++++++++++++++ agent/server/handler/failure_queue.go | 78 ---- agent/server/handler/handler_invoke.go | 74 +++- agent/server/handler/handler_manager_test.go | 11 +- agent/server/handler/manager.go | 205 ++++------- .../handler/scheduled_handler_entry_test.go | 21 +- agent/server/handler/source.go | 12 + agent/server/http/axon_handler_test.go | 4 +- agent/server/http/http_server_webhook_test.go | 12 +- agent/server/http/module.go | 1 + agent/server/server.go | 249 +++++++------ agent/server/test_axon_client.go | 20 +- .../cortexapps/axon/cortex-axon-agent.pb.go | 334 +++++++++++------- .../axon/cortex-axon-agent_grpc.pb.go | 85 +---- sdks/go/README.md | 8 +- sdks/go/axon_agent.go | 11 +- sdks/go/hooks.go | 99 +++--- sdks/go/hooks_test.go | 325 +++++++++-------- sdks/go/mock_axon/mocks.go | 30 +- sdks/python/README.md | 18 +- sdks/python/cortex_axon/axon_client.py | 81 ++--- sdks/python/cortex_axon/hooks.py | 26 +- sdks/python/example/main.py | 2 +- sdks/python/tests/test_hooks.py | 166 +++++---- 28 files changed, 1414 insertions(+), 1074 deletions(-) create mode 100644 agent/server/handler/failure_manager.go delete mode 100644 agent/server/handler/failure_queue.go create mode 100644 agent/server/handler/source.go diff --git a/agent/cmd/relay.go b/agent/cmd/relay.go index f3694629..0fd017d6 100644 --- a/agent/cmd/relay.go +++ b/agent/cmd/relay.go @@ -91,6 +91,7 @@ func buildRelayStack(cmd *cobra.Command, cfg config.AgentConfig, integrationInfo initStack(cmd, cfg, integrationInfo), AgentModule, fx.Provide(handler.NewHandlerManager), + fx.Provide(handler.NewFailureManager), relayModule, ) return stack diff --git a/agent/proto/cortex-axon-agent.proto b/agent/proto/cortex-axon-agent.proto index c4ca4266..11cc86c3 100644 --- a/agent/proto/cortex-axon-agent.proto +++ b/agent/proto/cortex-axon-agent.proto @@ -13,14 +13,12 @@ service AxonAgent { rpc Dispatch (stream DispatchRequest) returns (stream DispatchMessage); rpc ReportInvocation(ReportInvocationRequest) returns (ReportInvocationResponse); - // DispatchEvents streams failures the agent detected on its own: invocations - // that never reached the client, or that the client never reported. They use - // their own stream so a client that has stopped reading Dispatch because - // every handler is busy still hears about them. - rpc DispatchEvents(DispatchRequest) returns (stream DispatchEvent); - - // RetryInvocation queues a failed invocation again as a new attempt. - rpc RetryInvocation(RetryInvocationRequest) returns (RetryInvocationResponse); + // DispatchEvents sends the client every failed invocation, and the client + // answers each with what to do about it. It is its own stream so a client + // that has stopped reading Dispatch because every handler is busy still + // hears about failures. An event the client has not answered when the + // stream ends is sent again on the next one. + rpc DispatchEvents(stream DispatchEventsRequest) returns (stream DispatchEvent); } @@ -168,13 +166,26 @@ message DispatchEvent { InvocationFailure failure = 2; } -message RetryInvocationRequest { - DispatchHandlerInvoke invoke = 1; +enum FailureAction { + // Treated as DROP. + FAILURE_ACTION_UNSPECIFIED = 0; + // Do nothing more with the invocation. + FAILURE_ACTION_DROP = 1; + // Queue the invocation again as its next attempt. + FAILURE_ACTION_RETRY = 2; } -message RetryInvocationResponse { - Error error = 1; - string invocation_id = 2; +message FailureResolution { + string invocation_id = 1; + FailureAction action = 2; +} + +message DispatchEventsRequest { + // Identifies the client on the first message of the stream. + string dispatch_id = 1; + string client_version = 2; + // The client's answer to a DispatchEvent; unset on the first message. + FailureResolution resolution = 3; } // @@ -195,6 +206,10 @@ message ReportInvocationRequest { InvokeResult result = 200; Error error = 201; } + // Why the invocation failed, from a client that classifies its failures. + // The agent passes it on as a failure event; an error without it is only + // recorded. + InvocationFailure failure = 202; repeated Log logs = 300; } diff --git a/agent/server/failure_events_e2e_test.go b/agent/server/failure_events_e2e_test.go index a9b8d7f9..ae099d5b 100644 --- a/agent/server/failure_events_e2e_test.go +++ b/agent/server/failure_events_e2e_test.go @@ -5,7 +5,6 @@ import ( "errors" "net/url" "path/filepath" - "sync" "testing" "time" @@ -15,19 +14,16 @@ import ( "github.com/cortexapps/axon/server/handler" "github.com/stretchr/testify/require" "go.uber.org/zap" - "google.golang.org/protobuf/proto" + "google.golang.org/grpc" ) const failureTestWebhook = "failure-test-hook" type failureEventsFixture struct { - t *testing.T - ctx context.Context - manager handler.Manager - client *testAxonClient - - mu sync.Mutex - invokes []*pb.DispatchHandlerInvoke + t *testing.T + ctx context.Context + manager handler.Manager + client *testAxonClient received chan *pb.DispatchHandlerInvoke } @@ -49,7 +45,8 @@ func startFailureEventsFixture(t *testing.T, timeout time.Duration, respond func logger, _ := zap.NewDevelopment() manager := handler.NewHandlerManager(logger, cron.New(), nil) - agent := NewAxonAgent(Params{Logger: logger, Config: cfg, Manager: manager}) + failures := handler.NewFailureManager(logger, nil, manager) + agent := NewAxonAgent(Params{Logger: logger, Config: cfg, Manager: manager, Failures: failures}) require.NoError(t, agent.Start(context.Background())) t.Cleanup(agent.Close) @@ -108,25 +105,38 @@ func (f *failureEventsFixture) nextInvoke() *pb.DispatchHandlerInvoke { } } -func (f *failureEventsFixture) events() <-chan *pb.DispatchEvent { - stream, err := f.client.DispatchEvents(f.ctx) +type eventStream struct { + stream grpc.BidiStreamingClient[pb.DispatchEventsRequest, pb.DispatchEvent] + events chan *pb.DispatchEvent +} + +// events opens the failure event stream, as a client with a failure hook +// does. ctx bounds the stream, so cancelling it simulates a disconnect. +func (f *failureEventsFixture) events(ctx context.Context) *eventStream { + stream, err := f.client.DispatchEvents(ctx) require.NoError(f.t, err) - out := make(chan *pb.DispatchEvent, 10) + s := &eventStream{stream: stream, events: make(chan *pb.DispatchEvent, 10)} go func() { for { ev, err := stream.Recv() if err != nil { return } - out <- ev + s.events <- ev } }() - return out + return s } -func requireEvent(t *testing.T, events <-chan *pb.DispatchEvent) *pb.DispatchEvent { +func (s *eventStream) resolve(t *testing.T, ev *pb.DispatchEvent, action pb.FailureAction) { + require.NoError(t, s.stream.Send(&pb.DispatchEventsRequest{ + Resolution: &pb.FailureResolution{InvocationId: ev.Invoke.InvocationId, Action: action}, + })) +} + +func requireEvent(t *testing.T, s *eventStream) *pb.DispatchEvent { select { - case ev := <-events: + case ev := <-s.events: return ev case <-time.After(5 * time.Second): require.FailNow(t, "no failure event arrived") @@ -134,9 +144,9 @@ func requireEvent(t *testing.T, events <-chan *pb.DispatchEvent) *pb.DispatchEve } } -func requireNoEvent(t *testing.T, events <-chan *pb.DispatchEvent, wait time.Duration) { +func requireNoEvent(t *testing.T, s *eventStream, wait time.Duration) { select { - case ev := <-events: + case ev := <-s.events: require.FailNow(t, "unexpected failure event", "%v", ev) case <-time.After(wait): } @@ -158,19 +168,17 @@ func TestFailureEvents_UnreportedInvocationTimesOutAndRetries(t *testing.T) { require.Equal(t, int32(1), first.Attempt) time.Sleep(500 * time.Millisecond) - events := f.events() + events := f.events(f.ctx) ev := requireEvent(t, events) require.Equal(t, first.InvocationId, ev.Invoke.InvocationId) require.Equal(t, `{"id": 1}`, ev.Invoke.Args["body"]) require.Equal(t, pb.FailureReason_TIMEOUT, ev.Failure.Reason) require.Equal(t, pb.FailureSource_RUNTIME, ev.Failure.Source) - resp, err := f.client.RetryInvocation(f.ctx, ev.Invoke) - require.NoError(t, err) - require.Nil(t, resp.Error) + events.resolve(t, ev, pb.FailureAction_FAILURE_ACTION_RETRY) retried := f.nextInvoke() - require.Equal(t, resp.InvocationId, retried.InvocationId) + require.NotEqual(t, first.InvocationId, retried.InvocationId) require.Equal(t, int32(2), retried.Attempt) require.Equal(t, `{"id": 1}`, retried.Args["body"]) @@ -178,31 +186,49 @@ func TestFailureEvents_UnreportedInvocationTimesOutAndRetries(t *testing.T) { requireNoEvent(t, events, time.Second) } -func TestFailureEvents_ClientReportedFailureIsNotRepeated(t *testing.T) { +func TestFailureEvents_ClientReportedFailureIsDeliveredOnce(t *testing.T) { f := startFailureEventsFixture(t, 200*time.Millisecond, func(*pb.DispatchHandlerInvoke) (clientResult, error) { - // The client saw the failure and fires its own hook. - return clientResult{}, errors.New("handler blew up") + return clientResult{failure: &pb.InvocationFailure{ + Reason: pb.FailureReason_HANDLER_ERROR, + Source: pb.FailureSource_APPLICATION, + Message: "handler blew up", + }}, errors.New("handler blew up") }) - events := f.events() + events := f.events(f.ctx) f.postWebhook(`{}`) - f.nextInvoke() + invoke := f.nextInvoke() + + ev := requireEvent(t, events) + require.Equal(t, invoke.InvocationId, ev.Invoke.InvocationId) + require.Equal(t, pb.FailureReason_HANDLER_ERROR, ev.Failure.Reason) + require.Equal(t, pb.FailureSource_APPLICATION, ev.Failure.Source) + events.resolve(t, ev, pb.FailureAction_FAILURE_ACTION_DROP) - // Past the timeout and grace, the agent must stay quiet. + // Past the timeout and grace, the agent must not report it again, and + // dropping it must not run it again. requireNoEvent(t, events, time.Second) + select { + case again := <-f.received: + require.FailNow(t, "dropped invocation ran again", "%v", again) + default: + } } -func TestFailureEvents_RetryRejectedAtLimit(t *testing.T) { - f := startFailureEventsFixture(t, time.Second, func(*pb.DispatchHandlerInvoke) (clientResult, error) { - return clientResult{result: "ok"}, nil +func TestFailureEvents_UnresolvedEventIsResentOnReconnect(t *testing.T) { + f := startFailureEventsFixture(t, 200*time.Millisecond, func(*pb.DispatchHandlerInvoke) (clientResult, error) { + return clientResult{noReport: true}, nil }) f.postWebhook(`{}`) - invoke := proto.Clone(f.nextInvoke()).(*pb.DispatchHandlerInvoke) - invoke.Attempt = 5 + invoke := f.nextInvoke() - resp, err := f.client.RetryInvocation(f.ctx, invoke) - require.NoError(t, err) - require.NotNil(t, resp.Error) - require.Equal(t, "rejected", resp.Error.Code) + // The client receives the failure but disconnects before answering. + firstCtx, disconnect := context.WithCancel(f.ctx) + ev := requireEvent(t, f.events(firstCtx)) + require.Equal(t, invoke.InvocationId, ev.Invoke.InvocationId) + disconnect() + + again := requireEvent(t, f.events(f.ctx)) + require.Equal(t, invoke.InvocationId, again.Invoke.InvocationId) } diff --git a/agent/server/handler/failure_events_test.go b/agent/server/handler/failure_events_test.go index c04bdf1b..6e66278f 100644 --- a/agent/server/handler/failure_events_test.go +++ b/agent/server/handler/failure_events_test.go @@ -22,33 +22,44 @@ func webhookOnlyOption() *pb.HandlerOption { } } -func startWebhookHandler(t *testing.T, timeout time.Duration) (Manager, HandlerEntry) { +func startWebhookHandler(t *testing.T, timeout time.Duration) (Manager, FailureManager, HandlerEntry) { logger, _ := zap.NewDevelopment() mgr := NewHandlerManager(logger, cron.New(), nil) + failures := NewFailureManager(logger, nil, mgr) _, err := mgr.RegisterHandler("1", "handler1", timeout, webhookOnlyOption()) require.NoError(t, err) require.NoError(t, mgr.Start("1")) - return mgr, mgr.ListHandlers()[0] + return mgr, failures, mgr.ListHandlers()[0] } -func nextFailure(t *testing.T, mgr Manager) *pb.DispatchEvent { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - ev, err := mgr.NextFailure(ctx, "1") - require.NoError(t, err, "expected a failure event") +func nextFailure(t *testing.T, failures FailureManager) *pb.DispatchEvent { + ev, err := failures.Next(context.Background(), "1", 5*time.Second) + require.NoError(t, err) + require.NotNil(t, ev, "expected a failure event") return ev } -func requireNoFailure(t *testing.T, mgr Manager) { - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() - ev, err := mgr.NextFailure(ctx, "1") - require.Error(t, err, "unexpected failure event: %v", ev) +func requireNoFailure(t *testing.T, failures FailureManager) { + ev, err := failures.Next(context.Background(), "1", 200*time.Millisecond) + require.NoError(t, err) + require.Nil(t, ev, "unexpected failure event") } -func TestQueueFullEmitsNoResources(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) - mgr.SubscribeFailures("1") +func TestInvokeStateTransitions(t *testing.T) { + _, _, entry := startWebhookHandler(t, defaultTimeout) + invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil) + + require.Equal(t, InvokeQueued, invoke.State()) + require.True(t, invoke.SetState(InvokeDispatched)) + // It leaves the queue once: it can't then expire or be rejected. + require.False(t, invoke.SetState(InvokeExpired)) + require.False(t, invoke.SetState(InvokeRejected)) + require.Equal(t, InvokeDispatched, invoke.State()) +} + +func TestQueueFullIsAFailure(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") for i := 0; i < dispatchQueueDepth; i++ { require.NoError(t, mgr.Trigger(NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil))) @@ -56,8 +67,9 @@ func TestQueueFullEmitsNoResources(t *testing.T) { rejected := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"body": "payload"}) require.ErrorIs(t, mgr.Trigger(rejected), ErrDispatchQueueFull) + require.Equal(t, InvokeRejected, rejected.State()) - ev := nextFailure(t, mgr) + ev := nextFailure(t, failures) require.Equal(t, rejected.Id, ev.Invoke.InvocationId) require.Equal(t, "payload", ev.Invoke.Args["body"]) require.Equal(t, pb.FailureReason_NO_RESOURCES, ev.Failure.Reason) @@ -65,130 +77,190 @@ func TestQueueFullEmitsNoResources(t *testing.T) { require.True(t, ev.Failure.CallerRejected) } -func TestExpiredInQueueEmitsNoResources(t *testing.T) { - mgr, entry := startWebhookHandler(t, 50*time.Millisecond) - mgr.SubscribeFailures("1") +func TestExpiredInQueueIsAFailure(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, 50*time.Millisecond) + failures.Subscribe("1") invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil) require.NoError(t, mgr.Trigger(invoke)) - // Nothing dequeues, so it expires in the queue. - ev := nextFailure(t, mgr) + // Nothing takes it, so it expires in the queue. + ev := nextFailure(t, failures) require.Equal(t, invoke.Id, ev.Invoke.InvocationId) require.Equal(t, pb.FailureReason_NO_RESOURCES, ev.Failure.Reason) require.False(t, ev.Failure.CallerRejected) - // The dispatch loop must not send it once it has expired. - h, err := mgr.Dequeue(context.Background(), "1", time.Second) + // And it is never handed out afterwards. + h, err := mgr.Next(context.Background(), "1", 200*time.Millisecond) + require.NoError(t, err) + require.Nil(t, h) +} + +func TestDispatchedInvocationDoesNotExpire(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, 50*time.Millisecond) + failures.Subscribe("1") + + invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil) + require.NoError(t, mgr.Trigger(invoke)) + h, err := mgr.Next(context.Background(), "1", time.Second) require.NoError(t, err) - require.NotNil(t, h) - require.False(t, h.ClaimDispatch()) + require.Equal(t, InvokeDispatched, h.State()) + + // Once a client has it, only its report or the agent's timeout ends it. + time.Sleep(200 * time.Millisecond) + requireNoFailure(t, failures) + select { + case <-invoke.Done(): + require.Fail(t, "a dispatched invocation completed on its queue timeout") + default: + } } -func TestDispatchedInvocationDoesNotEmitOnExpiry(t *testing.T) { - mgr, entry := startWebhookHandler(t, 50*time.Millisecond) - mgr.SubscribeFailures("1") +func TestReportedFailureIsAFailure(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") - require.NoError(t, mgr.Trigger(NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil))) - h, err := mgr.Dequeue(context.Background(), "1", time.Second) + invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil) + require.NoError(t, mgr.Trigger(invoke)) + _, err := mgr.Next(context.Background(), "1", time.Second) require.NoError(t, err) - require.True(t, h.ClaimDispatch()) - // The agent's report timer owns timeouts for dispatched work. - <-h.Done() - requireNoFailure(t, mgr) + invoke.Complete("", NewFailureError(pb.FailureReason_HANDLER_ERROR, pb.FailureSource_APPLICATION, "boom")) + ev := nextFailure(t, failures) + require.Equal(t, invoke.Id, ev.Invoke.InvocationId) + require.Equal(t, pb.FailureReason_HANDLER_ERROR, ev.Failure.Reason) +} + +func TestPlainErrorIsNotAFailure(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") + + invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil) + require.NoError(t, mgr.Trigger(invoke)) + _, err := mgr.Next(context.Background(), "1", time.Second) + require.NoError(t, err) + + // An older client reports errors without a failure; that is only recorded. + invoke.Complete("", fmt.Errorf("invocation error: boom")) + requireNoFailure(t, failures) } func TestFailuresWithoutSubscriberAreNotQueued(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) + _, failures, entry := startWebhookHandler(t, defaultTimeout) - mgr.EmitFailure(NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke(), + failures.Notify(NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke(), &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) // Subscribing later must not surface the earlier event. - requireNoFailure(t, mgr) + requireNoFailure(t, failures) } func TestFailuresAreNeverDropped(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) - mgr.SubscribeFailures("1") + _, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") const count = 10000 for i := 0; i < count; i++ { invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"n": fmt.Sprint(i)}) - mgr.EmitFailure(invoke.ToDispatchInvoke(), &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) + failures.Notify(invoke.ToDispatchInvoke(), &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) } for i := 0; i < count; i++ { - ev := nextFailure(t, mgr) + ev := nextFailure(t, failures) require.Equal(t, fmt.Sprint(i), ev.Invoke.Args["n"], "events must arrive in order") } - requireNoFailure(t, mgr) + requireNoFailure(t, failures) } -func TestRequeuedFailureIsDeliveredNext(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) - mgr.SubscribeFailures("1") +func TestUnresolvedFailureIsRequeuedFirst(t *testing.T) { + _, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") for _, n := range []string{"a", "b"} { invoke := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"n": n}) - mgr.EmitFailure(invoke.ToDispatchInvoke(), &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) + failures.Notify(invoke.ToDispatchInvoke(), &pb.InvocationFailure{Reason: pb.FailureReason_TIMEOUT}) } - first := nextFailure(t, mgr) + first := nextFailure(t, failures) require.Equal(t, "a", first.Invoke.Args["n"]) - mgr.RequeueFailure("1", first) + failures.Requeue("1", first.Invoke.InvocationId) + + require.Equal(t, "a", nextFailure(t, failures).Invoke.Args["n"]) + require.Equal(t, "b", nextFailure(t, failures).Invoke.Args["n"]) +} + +func resolve(t *testing.T, failures FailureManager, ev *pb.DispatchEvent, action pb.FailureAction) error { + return failures.Resolve("1", &pb.FailureResolution{InvocationId: ev.Invoke.InvocationId, Action: action}) +} - require.Equal(t, "a", nextFailure(t, mgr).Invoke.Args["n"]) - require.Equal(t, "b", nextFailure(t, mgr).Invoke.Args["n"]) +func notifyFailure(failures FailureManager, invoke *pb.DispatchHandlerInvoke) { + failures.Notify(invoke, &pb.InvocationFailure{Reason: pb.FailureReason_HANDLER_ERROR}) } -func TestRetriggerQueuesNextAttempt(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) +func TestResolveRetryQueuesNextAttempt(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") original := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, map[string]string{"body": "payload"}) - id, err := mgr.Retrigger(original.ToDispatchInvoke()) - require.NoError(t, err) - require.NotEqual(t, original.Id, id) + notifyFailure(failures, original.ToDispatchInvoke()) + require.NoError(t, resolve(t, failures, nextFailure(t, failures), pb.FailureAction_FAILURE_ACTION_RETRY)) - h, err := mgr.Dequeue(context.Background(), "1", time.Second) + h, err := mgr.Next(context.Background(), "1", time.Second) require.NoError(t, err) retried := h.ToDispatchInvoke() - require.Equal(t, id, retried.InvocationId) + require.NotEqual(t, original.Id, retried.InvocationId) require.Equal(t, int32(2), retried.Attempt) require.Equal(t, "payload", retried.Args["body"]) require.Equal(t, pb.HandlerInvokeType_WEBHOOK, retried.Reason) } -func TestRetriggerFindsReregisteredHandler(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) +func TestResolveDropQueuesNothing(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") + + notifyFailure(failures, NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke()) + ev := nextFailure(t, failures) + require.NoError(t, resolve(t, failures, ev, pb.FailureAction_FAILURE_ACTION_DROP)) + + h, err := mgr.Next(context.Background(), "1", 200*time.Millisecond) + require.NoError(t, err) + require.Nil(t, h) + // Resolved once; a second answer finds nothing pending. + require.Error(t, resolve(t, failures, ev, pb.FailureAction_FAILURE_ACTION_RETRY)) +} + +func TestRetryFindsReregisteredHandler(t *testing.T) { + _, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") stale := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke() stale.HandlerId = "gone" + notifyFailure(failures, stale) - _, err := mgr.Retrigger(stale) - require.NoError(t, err) + require.NoError(t, resolve(t, failures, nextFailure(t, failures), pb.FailureAction_FAILURE_ACTION_RETRY)) } -func TestRetriggerStopsAtLimit(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) +func TestRetryStopsAtLimit(t *testing.T) { + _, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") last := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke() last.Attempt = maxRetryAttempts + notifyFailure(failures, last) - _, err := mgr.Retrigger(last) - require.ErrorIs(t, err, ErrRetryLimit) + require.ErrorIs(t, resolve(t, failures, nextFailure(t, failures), pb.FailureAction_FAILURE_ACTION_RETRY), ErrRetryLimit) } -func TestRetriggerTreatsMissingAttemptAsFirst(t *testing.T) { - mgr, entry := startWebhookHandler(t, defaultTimeout) +func TestRetryTreatsMissingAttemptAsFirst(t *testing.T) { + mgr, failures, entry := startWebhookHandler(t, defaultTimeout) + failures.Subscribe("1") fromOlderAgent := NewHandlerInvoke(entry, pb.HandlerInvokeType_WEBHOOK, nil).ToDispatchInvoke() fromOlderAgent.Attempt = 0 + notifyFailure(failures, fromOlderAgent) - _, err := mgr.Retrigger(fromOlderAgent) - require.NoError(t, err) - h, err := mgr.Dequeue(context.Background(), "1", time.Second) + require.NoError(t, resolve(t, failures, nextFailure(t, failures), pb.FailureAction_FAILURE_ACTION_RETRY)) + h, err := mgr.Next(context.Background(), "1", time.Second) require.NoError(t, err) require.Equal(t, int32(2), h.ToDispatchInvoke().Attempt) } diff --git a/agent/server/handler/failure_manager.go b/agent/server/handler/failure_manager.go new file mode 100644 index 00000000..d7554bfd --- /dev/null +++ b/agent/server/handler/failure_manager.go @@ -0,0 +1,258 @@ +package handler + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + pb "github.com/cortexapps/axon/.generated/proto/github.com/cortexapps/axon" + "github.com/prometheus/client_golang/prometheus" + "go.uber.org/zap" +) + +// maxRetryAttempts caps retries. The failure hook sees the attempt number +// and normally decides when to stop; this only bounds a hook that always +// retries. +const maxRetryAttempts = 5 + +// ErrRetryLimit means an invocation has used all of its attempts. +var ErrRetryLimit = errors.New("retry limit reached") + +// FailureNotifier receives every invocation that completed with a +// FailureError. +type FailureNotifier interface { + Notify(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) +} + +// FailureManager holds failed invocations until their client's failure hook +// says what to do with them: drop them, or retry them through the Manager. +type FailureManager interface { + FailureNotifier + // Next hands out the next failure for a client. It stays pending until + // Resolve or Requeue. + InvocationSource[*pb.DispatchEvent] + + // Subscribe starts keeping failures for a dispatch id. Failures for a + // dispatch id nobody subscribed are logged and discarded. + Subscribe(dispatchId string) + // Resolve applies the client's answer to a pending failure. + Resolve(dispatchId string, resolution *pb.FailureResolution) error + // Requeue puts a pending failure back to be sent again, first. + Requeue(dispatchId string, invocationId string) +} + +type failureManager struct { + logger *zap.Logger + manager Manager + + mu sync.Mutex + queues map[string]*failureQueue + pending prometheus.Gauge +} + +// NewFailureManager creates a FailureManager and has manager report its +// failures to it. +func NewFailureManager(logger *zap.Logger, metrics *prometheus.Registry, manager Manager) FailureManager { + fm := &failureManager{ + logger: logger.Named("failure-manager"), + manager: manager, + queues: make(map[string]*failureQueue), + pending: prometheus.NewGauge(prometheus.GaugeOpts{ + Name: "axon_handler_failure_events_pending", + Help: "Number of failure events waiting for a client to resolve them", + }), + } + if metrics != nil { + metrics.MustRegister(fm.pending) + } + manager.SetFailureNotifier(fm) + return fm +} + +func (f *failureManager) queue(dispatchId string, create bool) *failureQueue { + f.mu.Lock() + defer f.mu.Unlock() + q, ok := f.queues[dispatchId] + if !ok && create { + q = newFailureQueue() + f.queues[dispatchId] = q + } + return q +} + +func (f *failureManager) Subscribe(dispatchId string) { + f.queue(dispatchId, true) +} + +func (f *failureManager) Notify(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) { + logger := f.logger.With( + zap.String("handler", invoke.HandlerName), + zap.String("invocation-id", invoke.InvocationId), + zap.Int32("attempt", invoke.Attempt), + zap.String("reason", failure.Reason.String()), + zap.String("source", failure.Source.String()), + zap.String("message", failure.Message), + ) + q := f.queue(invoke.DispatchId, false) + if q == nil { + logger.Warn("Invocation failed, client does not handle failure events") + return + } + logger.Warn("Invocation failed, queueing failure event") + q.push(&pb.DispatchEvent{Invoke: invoke, Failure: failure}) + f.pending.Inc() +} + +func (f *failureManager) Next(ctx context.Context, dispatchId string, wait time.Duration) (*pb.DispatchEvent, error) { + return f.queue(dispatchId, true).next(ctx, wait), nil +} + +func (f *failureManager) Requeue(dispatchId string, invocationId string) { + f.queue(dispatchId, true).requeue(invocationId) +} + +func (f *failureManager) Resolve(dispatchId string, resolution *pb.FailureResolution) error { + ev := f.queue(dispatchId, true).take(resolution.GetInvocationId()) + if ev == nil { + return fmt.Errorf("no pending failure for invocation %s", resolution.GetInvocationId()) + } + f.pending.Dec() + + logger := f.logger.With( + zap.String("handler", ev.Invoke.HandlerName), + zap.String("invocation-id", ev.Invoke.InvocationId), + zap.String("action", resolution.GetAction().String()), + ) + if resolution.GetAction() != pb.FailureAction_FAILURE_ACTION_RETRY { + logger.Info("Dropping failed invocation") + return nil + } + id, err := f.retry(ev.Invoke) + if err != nil { + logger.Warn("Not retrying failed invocation", zap.Error(err)) + return err + } + logger.Info("Retrying failed invocation", zap.String("retry-invocation-id", id)) + return nil +} + +// retry queues invoke again as its next attempt. A retry that can't be +// queued fails like any other invocation, so it comes back as a failure. +func (f *failureManager) retry(invoke *pb.DispatchHandlerInvoke) (string, error) { + attempt := max(invoke.Attempt, 1) + if attempt >= maxRetryAttempts { + return "", fmt.Errorf("%w: attempt %d of %d", ErrRetryLimit, attempt, maxRetryAttempts) + } + + var entry HandlerEntry + for _, e := range f.manager.ListHandlers() { + if e.Id() == invoke.HandlerId { + entry = e + break + } + // The client re-registered since the invocation was sent, which + // gives the handler a new id under the same dispatch id and name. + if e.DispatchId() == invoke.DispatchId && e.Name() == invoke.HandlerName { + entry = e + } + } + if entry == nil { + return "", fmt.Errorf("handler not found: %s", invoke.HandlerName) + } + + next := NewHandlerInvoke(entry, invoke.Reason, invoke.Args) + next.Attempt = attempt + 1 + if err := f.manager.Trigger(next); err != nil { + return "", err + } + return next.Id, nil +} + +// failureQueue holds one dispatch id's failures until its client resolves +// them. It is unbounded on purpose: invocations are throttled, but a dropped +// failure is a message lost without a trace, which is what failure events +// exist to prevent. +type failureQueue struct { + mu sync.Mutex + events []*pb.DispatchEvent + pending map[string]*pb.DispatchEvent + // notify has room for one wakeup, so push never blocks and a waiting + // next never misses an event that arrives between checks. + notify chan struct{} +} + +func newFailureQueue() *failureQueue { + return &failureQueue{ + pending: make(map[string]*pb.DispatchEvent), + notify: make(chan struct{}, 1), + } +} + +func (q *failureQueue) push(ev *pb.DispatchEvent) { + q.mu.Lock() + q.events = append(q.events, ev) + q.mu.Unlock() + q.wake() +} + +func (q *failureQueue) wake() { + select { + case q.notify <- struct{}{}: + default: + } +} + +// next moves the oldest event to pending and returns it, waiting up to wait +// for one. It returns nil if none arrived. +func (q *failureQueue) next(ctx context.Context, wait time.Duration) *pb.DispatchEvent { + timer := time.NewTimer(wait) + defer timer.Stop() + for { + q.mu.Lock() + if len(q.events) > 0 { + ev := q.events[0] + q.events[0] = nil + q.events = q.events[1:] + q.pending[ev.Invoke.InvocationId] = ev + remaining := len(q.events) + q.mu.Unlock() + if remaining > 0 { + // Another stream may be waiting on the same queue. + q.wake() + } + return ev + } + q.mu.Unlock() + + select { + case <-q.notify: + case <-timer.C: + return nil + case <-ctx.Done(): + return nil + } + } +} + +func (q *failureQueue) take(invocationId string) *pb.DispatchEvent { + q.mu.Lock() + defer q.mu.Unlock() + ev := q.pending[invocationId] + delete(q.pending, invocationId) + return ev +} + +func (q *failureQueue) requeue(invocationId string) { + q.mu.Lock() + ev, ok := q.pending[invocationId] + if ok { + delete(q.pending, invocationId) + q.events = append([]*pb.DispatchEvent{ev}, q.events...) + } + q.mu.Unlock() + if ok { + q.wake() + } +} diff --git a/agent/server/handler/failure_queue.go b/agent/server/handler/failure_queue.go deleted file mode 100644 index c8b717b9..00000000 --- a/agent/server/handler/failure_queue.go +++ /dev/null @@ -1,78 +0,0 @@ -package handler - -import ( - "context" - "sync" - - pb "github.com/cortexapps/axon/.generated/proto/github.com/cortexapps/axon" -) - -// failureQueue holds failure events for one dispatch id until its client -// reads them. It is unbounded on purpose: invocations are throttled, but a -// dropped failure is a message lost without a trace, which is the problem -// failure events exist to solve. -type failureQueue struct { - mu sync.Mutex - events []*pb.DispatchEvent - // notify has room for one wakeup, so push never blocks and a waiting - // pop never misses an event that arrives while it is between checks. - notify chan struct{} -} - -func newFailureQueue() *failureQueue { - return &failureQueue{notify: make(chan struct{}, 1)} -} - -func (q *failureQueue) push(ev *pb.DispatchEvent) { - q.mu.Lock() - q.events = append(q.events, ev) - q.mu.Unlock() - q.wake() -} - -// pushFront returns an event whose delivery failed, so it is the next one sent. -func (q *failureQueue) pushFront(ev *pb.DispatchEvent) { - q.mu.Lock() - q.events = append([]*pb.DispatchEvent{ev}, q.events...) - q.mu.Unlock() - q.wake() -} - -func (q *failureQueue) wake() { - select { - case q.notify <- struct{}{}: - default: - } -} - -// pop blocks until an event is available or ctx is done. -func (q *failureQueue) pop(ctx context.Context) (*pb.DispatchEvent, error) { - for { - q.mu.Lock() - if len(q.events) > 0 { - ev := q.events[0] - q.events[0] = nil - q.events = q.events[1:] - remaining := len(q.events) - q.mu.Unlock() - if remaining > 0 { - // Another reader may be waiting on the same queue. - q.wake() - } - return ev, nil - } - q.mu.Unlock() - - select { - case <-q.notify: - case <-ctx.Done(): - return nil, ctx.Err() - } - } -} - -func (q *failureQueue) len() int { - q.mu.Lock() - defer q.mu.Unlock() - return len(q.events) -} diff --git a/agent/server/handler/handler_invoke.go b/agent/server/handler/handler_invoke.go index b93f9c9c..8ee53ac6 100644 --- a/agent/server/handler/handler_invoke.go +++ b/agent/server/handler/handler_invoke.go @@ -2,7 +2,9 @@ package handler import ( "context" + "errors" "fmt" + "slices" "sync/atomic" "time" @@ -20,18 +22,46 @@ type Invocable interface { Done() <-chan struct{} GetResult() (string, error) - // ClaimDispatch and ClaimExpired decide, exactly once, whether a queued - // invocation was sent to a client or expired before one picked it up. - ClaimDispatch() bool - ClaimExpired() bool + // SetState moves the invocation to a new state, and reports false when + // that transition isn't allowed from where it is now. + SetState(state InvokeState) bool + State() InvokeState } +type InvokeState int32 + const ( - invokeQueued int32 = iota - invokeDispatched - invokeExpired + // InvokeQueued is waiting for a client to take it. + InvokeQueued InvokeState = iota + // InvokeDispatched has been sent to a client. + InvokeDispatched + // InvokeExpired ran out of time before a client took it. + InvokeExpired + // InvokeRejected was refused because the queue was full. + InvokeRejected ) +// invokeTransitions lists the states each state may move to. A queued +// invocation leaves the queue exactly once, one way or another. +var invokeTransitions = map[InvokeState][]InvokeState{ + InvokeQueued: {InvokeDispatched, InvokeExpired, InvokeRejected}, +} + +// FailureError completes an invocation that failed in a way its handler's +// failure hook should hear about. Every such failure, wherever it is +// detected, reaches the failure manager this one way. +type FailureError struct { + Failure *pb.InvocationFailure +} + +func (e *FailureError) Error() string { + return fmt.Sprintf("%s (%s): %s", e.Failure.GetReason(), e.Failure.GetSource(), e.Failure.GetMessage()) +} + +func NewFailureError(reason pb.FailureReason, source pb.FailureSource, message string) *FailureError { + return &FailureError{Failure: &pb.InvocationFailure{Reason: reason, Source: source, Message: message}} +} + type HandlerInvoke struct { Id string Entry HandlerEntry @@ -82,14 +112,22 @@ func (h HandlerInvoke) GetReason() pb.HandlerInvokeType { return h.Reason } +// Start bounds how long the invocation may wait in the queue. Once a client +// has it, its result comes from the client's report or the agent's own +// timeout instead. func (h *HandlerInvoke) Start(ctx context.Context) { if h.started.CompareAndSwap(false, true) { go func() { select { case <-ctx.Done(): - err := ctx.Err() - if err == nil { - err = fmt.Errorf("context cancelled") + if !h.SetState(InvokeExpired) { + return + } + var err error = ctx.Err() + if errors.Is(err, context.DeadlineExceeded) { + err = NewFailureError(pb.FailureReason_NO_RESOURCES, pb.FailureSource_RUNTIME, + fmt.Sprintf("expired after %v waiting for a free handler", + time.Since(h.Timestamp).Round(time.Millisecond))) } h.Complete("", err) case <-h.done: @@ -122,12 +160,20 @@ func (h HandlerInvoke) Done() <-chan struct{} { return h.done } -func (h HandlerInvoke) ClaimDispatch() bool { - return h.state.CompareAndSwap(invokeQueued, invokeDispatched) +func (h HandlerInvoke) SetState(to InvokeState) bool { + for { + from := InvokeState(h.state.Load()) + if !slices.Contains(invokeTransitions[from], to) { + return false + } + if h.state.CompareAndSwap(int32(from), int32(to)) { + return true + } + } } -func (h HandlerInvoke) ClaimExpired() bool { - return h.state.CompareAndSwap(invokeQueued, invokeExpired) +func (h HandlerInvoke) State() InvokeState { + return InvokeState(h.state.Load()) } func (h HandlerInvoke) ToDispatchInvoke() *pb.DispatchHandlerInvoke { diff --git a/agent/server/handler/handler_manager_test.go b/agent/server/handler/handler_manager_test.go index b817b466..5dadacd0 100644 --- a/agent/server/handler/handler_manager_test.go +++ b/agent/server/handler/handler_manager_test.go @@ -100,7 +100,7 @@ func TestTriggerAndDequeue(t *testing.T) { invoke := NewScheduledHandlerInvoke(entry, pb.HandlerInvokeType_RUN_NOW) err = mgr.Trigger(invoke) require.NoError(t, err) - h, err := mgr.Dequeue(context.Background(), "1", 500*time.Millisecond) + h, err := mgr.Next(context.Background(), "1", 500*time.Millisecond) require.NoError(t, err) require.NotNil(t, h) require.Equal(t, "handler1", h.GetEntry().Name()) @@ -110,10 +110,9 @@ func TestTriggerAndDequeue(t *testing.T) { time.Sleep(time.Millisecond * 10) metrics, err := registry.Gather() require.NoError(t, err) - require.Equal(t, 4, len(metrics)) + require.Equal(t, 3, len(metrics)) expected := []string{ - "axon_handler_failure_events_pending", "axon_handler_invokes", "axon_handler_latency", "axon_handler_queue_depth", @@ -166,7 +165,7 @@ func TestTriggerAndDequeueCustomTimeout(t *testing.T) { err = mgr.Trigger(invoke) require.NoError(t, err) - h, err := mgr.Dequeue(context.Background(), "1", 500*time.Millisecond) + h, err := mgr.Next(context.Background(), "1", 500*time.Millisecond) require.NoError(t, err) require.NotNil(t, h) require.Equal(t, "handler1", h.GetEntry().Name()) @@ -181,7 +180,7 @@ func TestTriggerAndDequeueTimeout(t *testing.T) { _, err := mgr.RegisterHandler("1", "handler1", defaultTimeout, FixtureHandlerOption()) require.NoError(t, err) - h, err := mgr.Dequeue(context.Background(), "1", time.Millisecond) + h, err := mgr.Next(context.Background(), "1", time.Millisecond) require.NoError(t, err) require.Nil(t, h) } @@ -196,7 +195,7 @@ func TestTriggerAndDequeueTimeoutContext(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond) defer cancel() - h, err := mgr.Dequeue(ctx, "1", time.Hour) + h, err := mgr.Next(ctx, "1", time.Hour) require.NoError(t, err) require.Nil(t, h) } diff --git a/agent/server/handler/manager.go b/agent/server/handler/manager.go index 93bb69a1..ad038622 100644 --- a/agent/server/handler/manager.go +++ b/agent/server/handler/manager.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "os" - "sync" "time" pb "github.com/cortexapps/axon/.generated/proto/github.com/cortexapps/axon" @@ -23,14 +22,6 @@ const dispatchQueueDepth = 100 // response rather than accepting work that will only time out. var ErrDispatchQueueFull = errors.New("dispatch queue is full") -// maxRetryAttempts caps RetryInvocation. The failure hook sees the attempt -// number and normally decides when to stop; this only bounds a hook that -// always retries. -const maxRetryAttempts = 5 - -// ErrRetryLimit means an invocation has used all of its attempts. -var ErrRetryLimit = errors.New("retry limit reached") - type Manager interface { RegisterHandler(dispatchId string, name string, timeout time.Duration, options ...*pb.HandlerOption) (string, error) UnregisterHandler(id string) @@ -40,21 +31,15 @@ type Manager interface { Stop(id string) error Trigger(handler Invocable) error GetByTag(tag string) HandlerEntry - Dequeue(ctx context.Context, id string, waitTime time.Duration) (Invocable, error) + // Next hands out the next queued invocation, already moved to + // InvokeDispatched; ones that expired while queued are skipped. + InvocationSource[Invocable] Close() error IsFinished() bool - // SubscribeFailures starts keeping failure events for a dispatch id. - // Events for a dispatch id nobody subscribed are logged and discarded. - SubscribeFailures(dispatchId string) - EmitFailure(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) - // NextFailure blocks until a failure event is available or ctx is done. - NextFailure(ctx context.Context, dispatchId string) (*pb.DispatchEvent, error) - // RequeueFailure puts back an event that could not be delivered. - RequeueFailure(dispatchId string, ev *pb.DispatchEvent) - // Retrigger queues a failed invocation again as its next attempt and - // returns the new invocation id. - Retrigger(invoke *pb.DispatchHandlerInvoke) (string, error) + // SetFailureNotifier sets where invocations that complete with a + // FailureError are sent. + SetFailureNotifier(notifier FailureNotifier) } type handlerManager struct { @@ -68,10 +53,7 @@ type handlerManager struct { invokeCounter *prometheus.CounterVec queueDepthGauge *prometheus.GaugeVec handlerLatencyGauge *prometheus.HistogramVec - failuresPending prometheus.Gauge - - failureLock sync.Mutex - failureQueues map[string]*failureQueue + failures FailureNotifier } func NewHandlerManager(logger *zap.Logger, cron cron.Cron, metrics *prometheus.Registry) Manager { @@ -82,7 +64,6 @@ func NewHandlerManager(logger *zap.Logger, cron cron.Cron, metrics *prometheus.R dispatchQueues: make(map[string]chan Invocable), cron: cron, outstandingRequests: make(map[string]Invocable), - failureQueues: make(map[string]*failureQueue), invokeCounter: prometheus.NewCounterVec( prometheus.CounterOpts{ Name: "axon_handler_invokes", @@ -105,19 +86,12 @@ func NewHandlerManager(logger *zap.Logger, cron cron.Cron, metrics *prometheus.R }, []string{"handler", "result"}, ), - failuresPending: prometheus.NewGauge( - prometheus.GaugeOpts{ - Name: "axon_handler_failure_events_pending", - Help: "Number of failure events waiting for a client to read them", - }, - ), done: make(chan struct{}), } if metrics != nil { metrics.MustRegister(mgr.invokeCounter) metrics.MustRegister(mgr.queueDepthGauge) metrics.MustRegister(mgr.handlerLatencyGauge) - metrics.MustRegister(mgr.failuresPending) } return mgr } @@ -236,16 +210,28 @@ func (s *handlerManager) removeHandler(entry HandlerEntry) { delete(s.handlers, entry.Id()) } -func (s *handlerManager) Dequeue(ctx context.Context, dispatchId string, waitTime time.Duration) (Invocable, error) { +func (s *handlerManager) SetFailureNotifier(notifier FailureNotifier) { + s.failures = notifier +} + +func (s *handlerManager) Next(ctx context.Context, dispatchId string, waitTime time.Duration) (Invocable, error) { queue := s.getDispatchQueue(dispatchId) - select { - case <-time.After(waitTime): - s.checkFinished() - return nil, nil - case response := <-queue: - return response, nil - case <-ctx.Done(): - return nil, nil + timer := time.NewTimer(waitTime) + defer timer.Stop() + for { + select { + case <-timer.C: + s.checkFinished() + return nil, nil + case invoke := <-queue: + if invoke.SetState(InvokeDispatched) { + return invoke, nil + } + // It expired while queued, and that failure has been reported; + // handing it out now would run stale work. + case <-ctx.Done(): + return nil, nil + } } } @@ -280,43 +266,39 @@ func (s *handlerManager) Trigger(handler Invocable) error { entry.OnTrigger(message.GetReason()) - startTime := time.Now() queue := s.getDispatchQueue(entry.DispatchId()) + go s.watch(message, logger, cancel) select { case queue <- message: default: // Never block here: Trigger runs on the inbound HTTP goroutine for a // webhook, and a blocked send holds that request open with no timeout. - // Cancelling completes the invocation rather than leaking it. - cancel() - s.queueDepthGauge.WithLabelValues(handlerName).Set(float64(len(queue))) - s.invokeCounter.WithLabelValues(handlerName, "rejected").Inc() + message.SetState(InvokeRejected) + failure := NewFailureError(pb.FailureReason_NO_RESOURCES, pb.FailureSource_RUNTIME, + fmt.Sprintf("dispatch queue is full (%d waiting)", len(queue))) + failure.Failure.CallerRejected = true + message.Complete("", failure) logger.Warn("Dispatch queue is full, rejecting invocation", zap.Int("depth", len(queue))) - s.EmitFailure(message.ToDispatchInvoke(), &pb.InvocationFailure{ - Reason: pb.FailureReason_NO_RESOURCES, - Source: pb.FailureSource_RUNTIME, - Message: fmt.Sprintf("dispatch queue is full (%d waiting)", len(queue)), - CallerRejected: true, - }) return ErrDispatchQueueFull } s.queueDepthGauge.WithLabelValues(handlerName).Set(float64(len(queue))) + return nil +} - go func() { - defer cancel() - <-message.Done() - result, err := message.GetResult() - if errors.Is(err, context.DeadlineExceeded) && message.ClaimExpired() { - // No client picked it up in time, so it will never run. - s.EmitFailure(message.ToDispatchInvoke(), &pb.InvocationFailure{ - Reason: pb.FailureReason_NO_RESOURCES, - Source: pb.FailureSource_RUNTIME, - Message: fmt.Sprintf("expired after %v waiting for a free handler", - message.GetEntry().Timeout()), - }) - } - s.queueDepthGauge.WithLabelValues(handlerName).Set(float64(len(queue))) +// watch waits for an invocation to finish, however that happens, and is the +// one place a failed invocation is handed to the failure manager. +func (s *handlerManager) watch(invoke Invocable, logger *zap.Logger, cancel context.CancelFunc) { + defer cancel() + startTime := time.Now() + handlerName := invoke.GetEntry().Name() + <-invoke.Done() + result, err := invoke.GetResult() + + s.queueDepthGauge.WithLabelValues(handlerName).Set(float64(len(s.getDispatchQueue(invoke.GetEntry().DispatchId())))) + if invoke.State() == InvokeRejected { + s.invokeCounter.WithLabelValues(handlerName, "rejected").Inc() + } else { s.handlerLatencyGauge.WithLabelValues(handlerName, result).Observe(float64(time.Since(startTime).Milliseconds())) if err != nil { logger.Error("handler error", zap.Error(err)) @@ -325,85 +307,16 @@ func (s *handlerManager) Trigger(handler Invocable) error { logger.Info("Handler completed", zap.Int("result-length", len(result))) s.invokeCounter.WithLabelValues(handlerName, "success").Inc() } - }() - - return nil -} - -func (s *handlerManager) SubscribeFailures(dispatchId string) { - s.failureQueue(dispatchId, true) -} - -func (s *handlerManager) failureQueue(dispatchId string, create bool) *failureQueue { - s.failureLock.Lock() - defer s.failureLock.Unlock() - q, ok := s.failureQueues[dispatchId] - if !ok && create { - q = newFailureQueue() - s.failureQueues[dispatchId] = q - } - return q -} - -func (s *handlerManager) EmitFailure(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) { - logger := s.logger.With( - zap.String("handler", invoke.HandlerName), - zap.String("invocation-id", invoke.InvocationId), - zap.String("reason", failure.Reason.String()), - zap.String("source", failure.Source.String()), - zap.String("message", failure.Message), - ) - q := s.failureQueue(invoke.DispatchId, false) - if q == nil { - logger.Warn("Invocation failed, client does not handle failure events") - return - } - logger.Warn("Invocation failed, queueing failure event") - q.push(&pb.DispatchEvent{Invoke: invoke, Failure: failure}) - s.failuresPending.Inc() -} - -func (s *handlerManager) NextFailure(ctx context.Context, dispatchId string) (*pb.DispatchEvent, error) { - ev, err := s.failureQueue(dispatchId, true).pop(ctx) - if err != nil { - return nil, err - } - s.failuresPending.Dec() - return ev, nil -} - -func (s *handlerManager) RequeueFailure(dispatchId string, ev *pb.DispatchEvent) { - s.failureQueue(dispatchId, true).pushFront(ev) - s.failuresPending.Inc() -} - -func (s *handlerManager) Retrigger(invoke *pb.DispatchHandlerInvoke) (string, error) { - attempt := max(invoke.Attempt, 1) - if attempt >= maxRetryAttempts { - return "", fmt.Errorf("%w: attempt %d of %d", ErrRetryLimit, attempt, maxRetryAttempts) } - entry := s.handlers[invoke.HandlerId] - if entry == nil { - // The client re-registered since the invocation was sent, which gives - // the handler a new id under the same dispatch id and name. - for _, e := range s.handlers { - if e.DispatchId() == invoke.DispatchId && e.Name() == invoke.HandlerName { - entry = e - break - } + var failure *FailureError + if errors.As(err, &failure) { + if s.failures == nil { + logger.Warn("Invocation failed, no failure manager", zap.Error(err)) + return } + s.failures.Notify(invoke.ToDispatchInvoke(), failure.Failure) } - if entry == nil { - return "", fmt.Errorf("handler not found: %s", invoke.HandlerName) - } - - retry := NewHandlerInvoke(entry, invoke.Reason, invoke.Args) - retry.Attempt = attempt + 1 - if err := s.Trigger(retry); err != nil { - return "", err - } - return retry.Id, nil } func (s *handlerManager) getDispatchQueue(DispatchId string) chan Invocable { @@ -450,6 +363,10 @@ func TriggerInvoke(context context.Context, manager Manager, handlerName string, return "", err } - <-handlerInvoke.Done() - return handlerInvoke.GetResult() + select { + case <-handlerInvoke.Done(): + return handlerInvoke.GetResult() + case <-context.Done(): + return "", context.Err() + } } diff --git a/agent/server/handler/scheduled_handler_entry_test.go b/agent/server/handler/scheduled_handler_entry_test.go index 9f71f25f..a81186fe 100644 --- a/agent/server/handler/scheduled_handler_entry_test.go +++ b/agent/server/handler/scheduled_handler_entry_test.go @@ -253,7 +253,7 @@ func (fhm *fakeManager) Trigger(handler Invocable) error { return os.ErrNotExist } -func (fhm *fakeManager) Dequeue(ctx context.Context, id string, waitTime time.Duration) (Invocable, error) { +func (fhm *fakeManager) Next(ctx context.Context, id string, waitTime time.Duration) (Invocable, error) { panic("not implemented") // TODO: Implement } @@ -288,22 +288,5 @@ func (fc *fakeCron) Remove(id string) { delete(fc.handlers, id) } -func (fhm *fakeManager) SubscribeFailures(dispatchId string) { - panic("not implemented") // TODO: Implement -} - -func (fhm *fakeManager) EmitFailure(invoke *pb.DispatchHandlerInvoke, failure *pb.InvocationFailure) { - panic("not implemented") // TODO: Implement -} - -func (fhm *fakeManager) NextFailure(ctx context.Context, dispatchId string) (*pb.DispatchEvent, error) { - panic("not implemented") // TODO: Implement -} - -func (fhm *fakeManager) RequeueFailure(dispatchId string, ev *pb.DispatchEvent) { - panic("not implemented") // TODO: Implement -} - -func (fhm *fakeManager) Retrigger(invoke *pb.DispatchHandlerInvoke) (string, error) { - panic("not implemented") // TODO: Implement +func (fhm *fakeManager) SetFailureNotifier(notifier FailureNotifier) { } diff --git a/agent/server/handler/source.go b/agent/server/handler/source.go new file mode 100644 index 00000000..00523932 --- /dev/null +++ b/agent/server/handler/source.go @@ -0,0 +1,12 @@ +package handler + +import ( + "context" + "time" +) + +// InvocationSource hands out work for one dispatch id. Next waits up to wait +// and returns the zero value when nothing arrived in that time. +type InvocationSource[T comparable] interface { + Next(ctx context.Context, dispatchId string, wait time.Duration) (T, error) +} diff --git a/agent/server/http/axon_handler_test.go b/agent/server/http/axon_handler_test.go index 39060050..a928ff9d 100644 --- a/agent/server/http/axon_handler_test.go +++ b/agent/server/http/axon_handler_test.go @@ -53,7 +53,7 @@ func TestInvokeEndpoint(t *testing.T) { response := `{"status":"ok"}` go func() { - handlerInvocation, err := manager.Dequeue(context.Background(), "1", 500*time.Millisecond) + handlerInvocation, err := manager.Next(context.Background(), "1", 500*time.Millisecond) require.NoError(t, err) require.NotNil(t, handlerInvocation) require.Equal(t, "test-handler", handlerInvocation.GetEntry().Name()) @@ -111,7 +111,7 @@ func TestInvokeEndpointErr(t *testing.T) { payload := `{"body":"payload"}` go func() { - handlerInvocation, err := manager.Dequeue(context.Background(), "1", 500*time.Millisecond) + handlerInvocation, err := manager.Next(context.Background(), "1", 500*time.Millisecond) require.NoError(t, err) require.Equal(t, "test-handler", handlerInvocation.GetEntry().Name()) require.Equal(t, payload, handlerInvocation.ToDispatchInvoke().Args["body"]) diff --git a/agent/server/http/http_server_webhook_test.go b/agent/server/http/http_server_webhook_test.go index 28a2b560..a460fee5 100644 --- a/agent/server/http/http_server_webhook_test.go +++ b/agent/server/http/http_server_webhook_test.go @@ -32,7 +32,7 @@ func TestHandleWebhook(t *testing.T) { }, } - _, err := manager.RegisterHandler("1", "test", 1, option) + _, err := manager.RegisterHandler("1", "test", time.Minute, option) require.NoError(t, err) err = manager.Start("1") @@ -61,7 +61,7 @@ func TestHandleWebhook(t *testing.T) { Ensure that the handler was invoked */ assert.NotNil(t, manager.GetByTag("my-webhook-id").LastInvoked()) - handlerInvocation, err := manager.Dequeue(context.Background(), "1", 500*time.Millisecond) + handlerInvocation, err := manager.Next(context.Background(), "1", 500*time.Millisecond) require.NoError(t, err) require.Equal(t, "test", handlerInvocation.GetEntry().Name()) @@ -89,7 +89,8 @@ func TestHandleWebhookRejectsWhenQueueIsFull(t *testing.T) { _, err := manager.RegisterHandler("1", "test", time.Minute, option) require.NoError(t, err) require.NoError(t, manager.Start("1")) - manager.SubscribeFailures("1") + failures := handler.NewFailureManager(logger, nil, manager) + failures.Subscribe("1") webhookHandler := NewWebhookHandler(config.AgentConfig{}, logger, manager, nil) router := mux.NewRouter() @@ -125,10 +126,9 @@ func TestHandleWebhookRejectsWhenQueueIsFull(t *testing.T) { // The refused webhook still reaches the client's failure hook, flagged so // it can tell the sender was asked to retry. - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - ev, err := manager.NextFailure(ctx, "1") + ev, err := failures.Next(context.Background(), "1", 5*time.Second) require.NoError(t, err) + require.NotNil(t, ev) require.Equal(t, pb.FailureReason_NO_RESOURCES, ev.Failure.Reason) require.True(t, ev.Failure.CallerRejected) require.Equal(t, "payload", ev.Invoke.Args["body"]) diff --git a/agent/server/http/module.go b/agent/server/http/module.go index 531fab6f..e702b8a9 100644 --- a/agent/server/http/module.go +++ b/agent/server/http/module.go @@ -13,6 +13,7 @@ import ( var Module = fx.Module("handler", fx.Provide(handler.NewHandlerManager), + fx.Provide(handler.NewFailureManager), fx.Provide(cron.New), fx.Invoke(createWebhookHttpServer), ) diff --git a/agent/server/server.go b/agent/server/server.go index d8cd6627..d5ebc6e3 100644 --- a/agent/server/server.go +++ b/agent/server/server.go @@ -38,6 +38,7 @@ type AxonAgent struct { config config.AgentConfig logger *zap.Logger Manager handler.Manager + Failures handler.FailureManager cortexApiServer pb.CortexApiServer grpcServer *grpc.Server inflightLock sync.RWMutex @@ -59,7 +60,8 @@ type Params struct { Lifecycle fx.Lifecycle `optional:"true"` Logger *zap.Logger Config config.AgentConfig - Manager handler.Manager `optional:"true"` + Manager handler.Manager `optional:"true"` + Failures handler.FailureManager `optional:"true"` } func NewAxonAgent( @@ -70,6 +72,7 @@ func NewAxonAgent( config: p.Config, logger: logger, Manager: p.Manager, + Failures: p.Failures, cortexApiServer: api.NewCortexApiServer(logger, p.Config), outstandingRequests: make(map[string]inflightRequest), historyManager: handler.NewHistoryManager(p.Config, logger), @@ -139,8 +142,8 @@ func (s *AxonAgent) Dispatch(stream pb.AxonAgent_DispatchServer) error { // Subscribe before any handler starts, so a failure can't happen // before the client's DispatchEvents stream connects and be lost. - if req.FailureEvents { - s.Manager.SubscribeFailures(req.DispatchId) + if req.FailureEvents && s.Failures != nil { + s.Failures.Subscribe(req.DispatchId) } s.sendInvocations(req.DispatchId, stream) @@ -159,6 +162,28 @@ func (s *AxonAgent) sendComplete(stream pb.AxonAgent_DispatchServer) error { return nil } +// drain sends whatever src hands out for a dispatch id until ctx ends or a +// send fails. idle runs whenever src has nothing; returning true ends it. +func drain[T comparable](ctx context.Context, src handler.InvocationSource[T], dispatchId string, wait time.Duration, send func(T) error, idle func() bool) error { + var none T + for ctx.Err() == nil { + item, err := src.Next(ctx, dispatchId, wait) + if err != nil { + return err + } + if item == none { + if idle != nil && idle() { + return nil + } + continue + } + if err := send(item); err != nil { + return err + } + } + return nil +} + func (s *AxonAgent) sendInvocations(dispatchId string, stream pb.AxonAgent_DispatchServer) error { s.logger.Info("starting handler set", zap.String("dispatch-id", dispatchId)) @@ -168,88 +193,67 @@ func (s *AxonAgent) sendInvocations(dispatchId string, stream pb.AxonAgent_Dispa return err } - finished := false - go func() { - <-stream.Context().Done() - s.logger.Info("stream context closed", zap.String("dispatch-id", dispatchId), zap.Error(stream.Context().Err())) - finished = true - }() - go func() { defer func() { s.logger.Info("closing handler set", zap.String("dispatch-id", dispatchId)) s.Manager.Stop(dispatchId) }() - for !finished { - invoke, err := s.Manager.Dequeue(stream.Context(), dispatchId, s.config.DequeueWaitTime) - if err != nil { - s.logger.Error("failed to dequeue message", zap.Error(err)) - continue - } - - if invoke == nil { - - if s.Manager.IsFinished() && len(s.outstandingRequests) == 0 { - s.sendComplete(stream) - finished = true - } - continue - } - - if !invoke.ClaimDispatch() { - // It expired while queued and its failure was already - // reported; sending it now would run stale work. - continue - } - - now := time.Now() - msg := invoke.ToDispatchInvoke() - // The client may not read this for a while if every handler is - // busy, so it gets an absolute deadline rather than only a timeout. - deadline := now.Add(time.Duration(msg.TimeoutMs)*time.Millisecond + reportGrace) - msg.Deadline = timestamppb.New(deadline) - - dispatchMessage := &pb.DispatchMessage{ + send := func(invoke handler.Invocable) error { + return stream.Send(&pb.DispatchMessage{ Type: pb.DispatchMessageType_DISPATCH_MESSAGE_INVOKE, - Message: &pb.DispatchMessage_Invoke{Invoke: msg}, - } - - if err := stream.Send(dispatchMessage); err != nil { - s.logger.Error("failed to send dispatch message to client", zap.Error(err)) - return - } - - s.setOutstandingRequest(msg.InvocationId, &inflightRequest{ - invocable: invoke, - sentAt: now, + Message: &pb.DispatchMessage_Invoke{Invoke: s.track(invoke)}, }) - go func(requestId string) { - // The grace lets a client that enforces the timeout itself - // report it first, so the failure is described once, by the - // side that saw it. - <-time.After(time.Until(deadline)) - message := fmt.Sprintf("handler did not report a result within %dms", msg.TimeoutMs) - completed := s.completeInvocation(context.Background(), &pb.ReportInvocationRequest{ - HandlerInvoke: msg, - Message: &pb.ReportInvocationRequest_Error{Error: &pb.Error{ - Code: "timeout", - Message: message, - }}, - }) - if completed { - s.Manager.EmitFailure(msg, &pb.InvocationFailure{ - Reason: pb.FailureReason_TIMEOUT, - Source: pb.FailureSource_RUNTIME, - Message: message, - }) - } - }(msg.InvocationId) + } + idle := func() bool { + if s.Manager.IsFinished() && s.outstandingCount() == 0 { + s.sendComplete(stream) + return true + } + return false + } + err := drain(stream.Context(), s.Manager, dispatchId, s.config.DequeueWaitTime, send, idle) + if err != nil { + s.logger.Error("failed to send dispatch message to client", zap.Error(err)) } }() return nil } +// track records an invocation as sent and gives it a deadline. If no report +// comes back by then, the agent completes it as a timeout, which it reports +// like any other failure. The grace lets a client that enforces the timeout +// itself report it first, so it is described by the side that saw it. +func (s *AxonAgent) track(invoke handler.Invocable) *pb.DispatchHandlerInvoke { + now := time.Now() + msg := invoke.ToDispatchInvoke() + // The client may not read this for a while if every handler is busy, so + // it gets an absolute deadline rather than only a timeout. + deadline := now.Add(time.Duration(msg.TimeoutMs)*time.Millisecond + reportGrace) + msg.Deadline = timestamppb.New(deadline) + + s.setOutstandingRequest(msg.InvocationId, &inflightRequest{invocable: invoke, sentAt: now}) + time.AfterFunc(time.Until(deadline), func() { + message := fmt.Sprintf("handler did not report a result within %dms", msg.TimeoutMs) + s.completeInvocation(context.Background(), &pb.ReportInvocationRequest{ + HandlerInvoke: msg, + Message: &pb.ReportInvocationRequest_Error{Error: &pb.Error{Code: "timeout", Message: message}}, + Failure: &pb.InvocationFailure{ + Reason: pb.FailureReason_TIMEOUT, + Source: pb.FailureSource_RUNTIME, + Message: message, + }, + }) + }) + return msg +} + +func (s *AxonAgent) outstandingCount() int { + s.inflightLock.RLock() + defer s.inflightLock.RUnlock() + return len(s.outstandingRequests) +} + func (s *AxonAgent) setOutstandingRequest(id string, req *inflightRequest) { s.inflightLock.Lock() defer s.inflightLock.Unlock() @@ -285,16 +289,18 @@ func (s *AxonAgent) ReportInvocation(ctx context.Context, req *pb.ReportInvocati return &pb.ReportInvocationResponse{}, nil } -// completeInvocation records the result of an outstanding invocation. It -// returns false when the invocation was already completed by an earlier report. -func (s *AxonAgent) completeInvocation(ctx context.Context, req *pb.ReportInvocationRequest) bool { +// completeInvocation records the result of an outstanding invocation, once: +// a later report for the same invocation is ignored. A report that carries a +// failure completes the invocation with a FailureError, which is how both +// client and agent failures reach the failure manager. +func (s *AxonAgent) completeInvocation(ctx context.Context, req *pb.ReportInvocationRequest) { s.inflightLock.Lock() ifr, ok := s.outstandingRequests[req.HandlerInvoke.InvocationId] delete(s.outstandingRequests, req.HandlerInvoke.InvocationId) s.inflightLock.Unlock() if !ok { - return false + return } var requestErr error @@ -302,6 +308,9 @@ func (s *AxonAgent) completeInvocation(ctx context.Context, req *pb.ReportInvoca if e := req.GetError(); e != nil { requestErr = fmt.Errorf("invocation error: %s", describeError(e)) } + if f := req.GetFailure(); f != nil { + requestErr = &handler.FailureError{Failure: f} + } requestResult := "" if rr := req.GetResult(); rr != nil { @@ -314,50 +323,68 @@ func (s *AxonAgent) completeInvocation(ctx context.Context, req *pb.ReportInvoca if err != nil { s.logger.Error("failed to write history file", zap.Error(err)) } - return true } -// DispatchEvents streams the agent's failure events for a dispatch id. An -// event whose send fails goes back on the queue for the next stream. -func (s *AxonAgent) DispatchEvents(req *pb.DispatchRequest, stream pb.AxonAgent_DispatchEventsServer) error { - if s.Manager == nil { - return fmt.Errorf("handler manager is not initialized") +// DispatchEvents sends a client its failed invocations and applies its answer +// to each. One the client hasn't answered when the stream ends is sent again +// on its next stream. +func (s *AxonAgent) DispatchEvents(stream pb.AxonAgent_DispatchEventsServer) error { + if s.Failures == nil { + return status.Error(codes.Unimplemented, "failure events are not enabled") } - s.logger.Info("client subscribed to failure events", zap.String("dispatch-id", req.DispatchId)) - s.Manager.SubscribeFailures(req.DispatchId) - for { - ev, err := s.Manager.NextFailure(stream.Context(), req.DispatchId) - if err != nil { - return nil + first, err := stream.Recv() + if err != nil { + return err + } + dispatchId := first.GetDispatchId() + logger := s.logger.With(zap.String("dispatch-id", dispatchId)) + logger.Info("client subscribed to failure events", zap.String("client-version", first.GetClientVersion())) + s.Failures.Subscribe(dispatchId) + + var mu sync.Mutex + unanswered := map[string]bool{} + defer func() { + mu.Lock() + defer mu.Unlock() + for id := range unanswered { + s.Failures.Requeue(dispatchId, id) } - if err := stream.Send(ev); err != nil { - s.Manager.RequeueFailure(req.DispatchId, ev) - s.logger.Warn("failed to send failure event, requeued", - zap.String("dispatch-id", req.DispatchId), zap.Error(err)) - return err + }() + + ctx, cancel := context.WithCancel(stream.Context()) + defer cancel() + go func() { + defer cancel() + for { + req, err := stream.Recv() + if err != nil { + return + } + res := req.GetResolution() + if res == nil { + continue + } + mu.Lock() + delete(unanswered, res.GetInvocationId()) + mu.Unlock() + if err := s.Failures.Resolve(dispatchId, res); err != nil { + logger.Warn("failed to apply failure resolution", zap.Error(err)) + } } - } -} + }() -// RetryInvocation queues a failed invocation again as its next attempt. -func (s *AxonAgent) RetryInvocation(ctx context.Context, req *pb.RetryInvocationRequest) (*pb.RetryInvocationResponse, error) { - if s.Manager == nil { - return nil, fmt.Errorf("handler manager is not initialized") - } - if req.Invoke == nil { - return &pb.RetryInvocationResponse{Error: &pb.Error{Code: "invalid", Message: "invoke is required"}}, nil + send := func(ev *pb.DispatchEvent) error { + mu.Lock() + unanswered[ev.Invoke.InvocationId] = true + mu.Unlock() + return stream.Send(ev) } - - id, err := s.Manager.Retrigger(req.Invoke) - if err != nil { - s.logger.Warn("retry rejected", - zap.String("handler", req.Invoke.HandlerName), - zap.String("invocation-id", req.Invoke.InvocationId), - zap.Error(err)) - return &pb.RetryInvocationResponse{Error: &pb.Error{Code: "rejected", Message: err.Error()}}, nil + if err := drain(ctx, s.Failures, dispatchId, s.config.DequeueWaitTime, send, nil); err != nil { + logger.Warn("failure event stream ended", zap.Error(err)) + return err } - return &pb.RetryInvocationResponse{InvocationId: id}, nil + return nil } // ListHandlers returns a list of all registered handlers diff --git a/agent/server/test_axon_client.go b/agent/server/test_axon_client.go index d894abb0..db2d364f 100644 --- a/agent/server/test_axon_client.go +++ b/agent/server/test_axon_client.go @@ -23,6 +23,9 @@ type clientResult struct { result string // noReport skips ReportInvocation, like a client that hung. noReport bool + // failure is reported with an error, as a client that classifies its + // failures does. + failure *pb.InvocationFailure } type testAxonClient struct { @@ -122,6 +125,7 @@ func (c *testAxonClient) Run(ctx context.Context) error { Message: err.Error(), }, } + report.Failure = result.failure } else { report.Message = &pb.ReportInvocationRequest_Result{ Result: &pb.InvokeResult{ @@ -162,10 +166,14 @@ func (c *testAxonClient) GetHandlerHistory(ctx context.Context, name string) ([] return resp.History, nil } -func (c *testAxonClient) DispatchEvents(ctx context.Context) (grpc.ServerStreamingClient[pb.DispatchEvent], error) { - return c.client.DispatchEvents(ctx, &pb.DispatchRequest{DispatchId: c.dispatchId, FailureEvents: true}) -} - -func (c *testAxonClient) RetryInvocation(ctx context.Context, invoke *pb.DispatchHandlerInvoke) (*pb.RetryInvocationResponse, error) { - return c.client.RetryInvocation(ctx, &pb.RetryInvocationRequest{Invoke: invoke}) +// DispatchEvents opens the failure event stream and sends the first message. +func (c *testAxonClient) DispatchEvents(ctx context.Context) (grpc.BidiStreamingClient[pb.DispatchEventsRequest, pb.DispatchEvent], error) { + stream, err := c.client.DispatchEvents(ctx) + if err != nil { + return nil, err + } + if err := stream.Send(&pb.DispatchEventsRequest{DispatchId: c.dispatchId}); err != nil { + return nil, err + } + return stream, nil } diff --git a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go index 36684fc1..15bb5669 100644 --- a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go +++ b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent.pb.go @@ -233,6 +233,58 @@ func (FailureSource) EnumDescriptor() ([]byte, []int) { return file_cortex_axon_agent_proto_rawDescGZIP(), []int{3} } +type FailureAction int32 + +const ( + // Treated as DROP. + FailureAction_FAILURE_ACTION_UNSPECIFIED FailureAction = 0 + // Do nothing more with the invocation. + FailureAction_FAILURE_ACTION_DROP FailureAction = 1 + // Queue the invocation again as its next attempt. + FailureAction_FAILURE_ACTION_RETRY FailureAction = 2 +) + +// Enum value maps for FailureAction. +var ( + FailureAction_name = map[int32]string{ + 0: "FAILURE_ACTION_UNSPECIFIED", + 1: "FAILURE_ACTION_DROP", + 2: "FAILURE_ACTION_RETRY", + } + FailureAction_value = map[string]int32{ + "FAILURE_ACTION_UNSPECIFIED": 0, + "FAILURE_ACTION_DROP": 1, + "FAILURE_ACTION_RETRY": 2, + } +) + +func (x FailureAction) Enum() *FailureAction { + p := new(FailureAction) + *p = x + return p +} + +func (x FailureAction) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (FailureAction) Descriptor() protoreflect.EnumDescriptor { + return file_cortex_axon_agent_proto_enumTypes[4].Descriptor() +} + +func (FailureAction) Type() protoreflect.EnumType { + return &file_cortex_axon_agent_proto_enumTypes[4] +} + +func (x FailureAction) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use FailureAction.Descriptor instead. +func (FailureAction) EnumDescriptor() ([]byte, []int) { + return file_cortex_axon_agent_proto_rawDescGZIP(), []int{4} +} + type RegisterHandlerRequest struct { state protoimpl.MessageState `protogen:"open.v1"` DispatchId string `protobuf:"bytes,1,opt,name=dispatch_id,json=dispatchId,proto3" json:"dispatch_id,omitempty"` @@ -1101,27 +1153,28 @@ func (x *DispatchEvent) GetFailure() *InvocationFailure { return nil } -type RetryInvocationRequest struct { +type FailureResolution struct { state protoimpl.MessageState `protogen:"open.v1"` - Invoke *DispatchHandlerInvoke `protobuf:"bytes,1,opt,name=invoke,proto3" json:"invoke,omitempty"` + InvocationId string `protobuf:"bytes,1,opt,name=invocation_id,json=invocationId,proto3" json:"invocation_id,omitempty"` + Action FailureAction `protobuf:"varint,2,opt,name=action,proto3,enum=cortex.axon.FailureAction" json:"action,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } -func (x *RetryInvocationRequest) Reset() { - *x = RetryInvocationRequest{} +func (x *FailureResolution) Reset() { + *x = FailureResolution{} mi := &file_cortex_axon_agent_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } -func (x *RetryInvocationRequest) String() string { +func (x *FailureResolution) String() string { return protoimpl.X.MessageStringOf(x) } -func (*RetryInvocationRequest) ProtoMessage() {} +func (*FailureResolution) ProtoMessage() {} -func (x *RetryInvocationRequest) ProtoReflect() protoreflect.Message { +func (x *FailureResolution) ProtoReflect() protoreflect.Message { mi := &file_cortex_axon_agent_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -1133,40 +1186,50 @@ func (x *RetryInvocationRequest) ProtoReflect() protoreflect.Message { return mi.MessageOf(x) } -// Deprecated: Use RetryInvocationRequest.ProtoReflect.Descriptor instead. -func (*RetryInvocationRequest) Descriptor() ([]byte, []int) { +// Deprecated: Use FailureResolution.ProtoReflect.Descriptor instead. +func (*FailureResolution) Descriptor() ([]byte, []int) { return file_cortex_axon_agent_proto_rawDescGZIP(), []int{14} } -func (x *RetryInvocationRequest) GetInvoke() *DispatchHandlerInvoke { +func (x *FailureResolution) GetInvocationId() string { if x != nil { - return x.Invoke + return x.InvocationId } - return nil + return "" } -type RetryInvocationResponse struct { - state protoimpl.MessageState `protogen:"open.v1"` - Error *Error `protobuf:"bytes,1,opt,name=error,proto3" json:"error,omitempty"` - InvocationId string `protobuf:"bytes,2,opt,name=invocation_id,json=invocationId,proto3" json:"invocation_id,omitempty"` +func (x *FailureResolution) GetAction() FailureAction { + if x != nil { + return x.Action + } + return FailureAction_FAILURE_ACTION_UNSPECIFIED +} + +type DispatchEventsRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Identifies the client on the first message of the stream. + DispatchId string `protobuf:"bytes,1,opt,name=dispatch_id,json=dispatchId,proto3" json:"dispatch_id,omitempty"` + ClientVersion string `protobuf:"bytes,2,opt,name=client_version,json=clientVersion,proto3" json:"client_version,omitempty"` + // The client's answer to a DispatchEvent; unset on the first message. + Resolution *FailureResolution `protobuf:"bytes,3,opt,name=resolution,proto3" json:"resolution,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } -func (x *RetryInvocationResponse) Reset() { - *x = RetryInvocationResponse{} +func (x *DispatchEventsRequest) Reset() { + *x = DispatchEventsRequest{} mi := &file_cortex_axon_agent_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } -func (x *RetryInvocationResponse) String() string { +func (x *DispatchEventsRequest) String() string { return protoimpl.X.MessageStringOf(x) } -func (*RetryInvocationResponse) ProtoMessage() {} +func (*DispatchEventsRequest) ProtoMessage() {} -func (x *RetryInvocationResponse) ProtoReflect() protoreflect.Message { +func (x *DispatchEventsRequest) ProtoReflect() protoreflect.Message { mi := &file_cortex_axon_agent_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) @@ -1178,25 +1241,32 @@ func (x *RetryInvocationResponse) ProtoReflect() protoreflect.Message { return mi.MessageOf(x) } -// Deprecated: Use RetryInvocationResponse.ProtoReflect.Descriptor instead. -func (*RetryInvocationResponse) Descriptor() ([]byte, []int) { +// Deprecated: Use DispatchEventsRequest.ProtoReflect.Descriptor instead. +func (*DispatchEventsRequest) Descriptor() ([]byte, []int) { return file_cortex_axon_agent_proto_rawDescGZIP(), []int{15} } -func (x *RetryInvocationResponse) GetError() *Error { +func (x *DispatchEventsRequest) GetDispatchId() string { if x != nil { - return x.Error + return x.DispatchId } - return nil + return "" } -func (x *RetryInvocationResponse) GetInvocationId() string { +func (x *DispatchEventsRequest) GetClientVersion() string { if x != nil { - return x.InvocationId + return x.ClientVersion } return "" } +func (x *DispatchEventsRequest) GetResolution() *FailureResolution { + if x != nil { + return x.Resolution + } + return nil +} + type Log struct { state protoimpl.MessageState `protogen:"open.v1"` Level string `protobuf:"bytes,1,opt,name=level,proto3" json:"level,omitempty"` @@ -1266,8 +1336,12 @@ type ReportInvocationRequest struct { // // *ReportInvocationRequest_Result // *ReportInvocationRequest_Error - Message isReportInvocationRequest_Message `protobuf_oneof:"message"` - Logs []*Log `protobuf:"bytes,300,rep,name=logs,proto3" json:"logs,omitempty"` + Message isReportInvocationRequest_Message `protobuf_oneof:"message"` + // Why the invocation failed, from a client that classifies its failures. + // The agent passes it on as a failure event; an error without it is only + // recorded. + Failure *InvocationFailure `protobuf:"bytes,202,opt,name=failure,proto3" json:"failure,omitempty"` + Logs []*Log `protobuf:"bytes,300,rep,name=logs,proto3" json:"logs,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -1348,6 +1422,13 @@ func (x *ReportInvocationRequest) GetError() *Error { return nil } +func (x *ReportInvocationRequest) GetFailure() *InvocationFailure { + if x != nil { + return x.Failure + } + return nil +} + func (x *ReportInvocationRequest) GetLogs() []*Log { if x != nil { return x.Logs @@ -1774,23 +1855,29 @@ const file_cortex_axon_agent_proto_rawDesc = "" + "\x0fcaller_rejected\x18\x04 \x01(\bR\x0ecallerRejected\"\x85\x01\n" + "\rDispatchEvent\x12:\n" + "\x06invoke\x18\x01 \x01(\v2\".cortex.axon.DispatchHandlerInvokeR\x06invoke\x128\n" + - "\afailure\x18\x02 \x01(\v2\x1e.cortex.axon.InvocationFailureR\afailure\"T\n" + - "\x16RetryInvocationRequest\x12:\n" + - "\x06invoke\x18\x01 \x01(\v2\".cortex.axon.DispatchHandlerInvokeR\x06invoke\"h\n" + - "\x17RetryInvocationResponse\x12(\n" + - "\x05error\x18\x01 \x01(\v2\x12.cortex.axon.ErrorR\x05error\x12#\n" + - "\rinvocation_id\x18\x02 \x01(\tR\finvocationId\"o\n" + + "\afailure\x18\x02 \x01(\v2\x1e.cortex.axon.InvocationFailureR\afailure\"l\n" + + "\x11FailureResolution\x12#\n" + + "\rinvocation_id\x18\x01 \x01(\tR\finvocationId\x122\n" + + "\x06action\x18\x02 \x01(\x0e2\x1a.cortex.axon.FailureActionR\x06action\"\x9f\x01\n" + + "\x15DispatchEventsRequest\x12\x1f\n" + + "\vdispatch_id\x18\x01 \x01(\tR\n" + + "dispatchId\x12%\n" + + "\x0eclient_version\x18\x02 \x01(\tR\rclientVersion\x12>\n" + + "\n" + + "resolution\x18\x03 \x01(\v2\x1e.cortex.axon.FailureResolutionR\n" + + "resolution\"o\n" + "\x03Log\x12\x14\n" + "\x05level\x18\x01 \x01(\tR\x05level\x128\n" + "\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12\x18\n" + - "\amessage\x18\x03 \x01(\tR\amessage\"\xec\x02\n" + + "\amessage\x18\x03 \x01(\tR\amessage\"\xa7\x03\n" + "\x17ReportInvocationRequest\x12I\n" + "\x0ehandler_invoke\x18\x01 \x01(\v2\".cortex.axon.DispatchHandlerInvokeR\rhandlerInvoke\x12P\n" + "\x16start_client_timestamp\x18d \x01(\v2\x1a.google.protobuf.TimestampR\x14startClientTimestamp\x12\x1f\n" + "\vduration_ms\x18e \x01(\x05R\n" + "durationMs\x124\n" + "\x06result\x18\xc8\x01 \x01(\v2\x19.cortex.axon.InvokeResultH\x00R\x06result\x12+\n" + - "\x05error\x18\xc9\x01 \x01(\v2\x12.cortex.axon.ErrorH\x00R\x05error\x12%\n" + + "\x05error\x18\xc9\x01 \x01(\v2\x12.cortex.axon.ErrorH\x00R\x05error\x129\n" + + "\afailure\x18\xca\x01 \x01(\v2\x1e.cortex.axon.InvocationFailureR\afailure\x12%\n" + "\x04logs\x18\xac\x02 \x03(\v2\x10.cortex.axon.LogR\x04logsB\t\n" + "\amessage\"D\n" + "\x18ReportInvocationResponse\x12(\n" + @@ -1842,16 +1929,19 @@ const file_cortex_axon_agent_proto_rawDesc = "" + "\rFailureSource\x12\x1e\n" + "\x1aFAILURE_SOURCE_UNSPECIFIED\x10\x00\x12\x0f\n" + "\vAPPLICATION\x10\x01\x12\v\n" + - "\aRUNTIME\x10\x022\xdf\x05\n" + + "\aRUNTIME\x10\x02*b\n" + + "\rFailureAction\x12\x1e\n" + + "\x1aFAILURE_ACTION_UNSPECIFIED\x10\x00\x12\x17\n" + + "\x13FAILURE_ACTION_DROP\x10\x01\x12\x18\n" + + "\x14FAILURE_ACTION_RETRY\x10\x022\x89\x05\n" + "\tAxonAgent\x12\\\n" + "\x0fRegisterHandler\x12#.cortex.axon.RegisterHandlerRequest\x1a$.cortex.axon.RegisterHandlerResponse\x12b\n" + "\x11UnregisterHandler\x12%.cortex.axon.UnregisterHandlerRequest\x1a&.cortex.axon.UnregisterHandlerResponse\x12S\n" + "\fListHandlers\x12 .cortex.axon.ListHandlersRequest\x1a!.cortex.axon.ListHandlersResponse\x12b\n" + "\x11GetHandlerHistory\x12%.cortex.axon.GetHandlerHistoryRequest\x1a&.cortex.axon.GetHandlerHistoryResponse\x12J\n" + "\bDispatch\x12\x1c.cortex.axon.DispatchRequest\x1a\x1c.cortex.axon.DispatchMessage(\x010\x01\x12_\n" + - "\x10ReportInvocation\x12$.cortex.axon.ReportInvocationRequest\x1a%.cortex.axon.ReportInvocationResponse\x12L\n" + - "\x0eDispatchEvents\x12\x1c.cortex.axon.DispatchRequest\x1a\x1a.cortex.axon.DispatchEvent0\x01\x12\\\n" + - "\x0fRetryInvocation\x12#.cortex.axon.RetryInvocationRequest\x1a$.cortex.axon.RetryInvocationResponseB\x1cZ\x1agithub.com/cortexapps/axonb\x06proto3" + "\x10ReportInvocation\x12$.cortex.axon.ReportInvocationRequest\x1a%.cortex.axon.ReportInvocationResponse\x12T\n" + + "\x0eDispatchEvents\x12\".cortex.axon.DispatchEventsRequest\x1a\x1a.cortex.axon.DispatchEvent(\x010\x01B\x1cZ\x1agithub.com/cortexapps/axonb\x06proto3" var ( file_cortex_axon_agent_proto_rawDescOnce sync.Once @@ -1865,98 +1955,98 @@ func file_cortex_axon_agent_proto_rawDescGZIP() []byte { return file_cortex_axon_agent_proto_rawDescData } -var file_cortex_axon_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 4) +var file_cortex_axon_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 5) var file_cortex_axon_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 24) var file_cortex_axon_agent_proto_goTypes = []any{ (HandlerInvokeType)(0), // 0: cortex.axon.HandlerInvokeType (DispatchMessageType)(0), // 1: cortex.axon.DispatchMessageType (FailureReason)(0), // 2: cortex.axon.FailureReason (FailureSource)(0), // 3: cortex.axon.FailureSource - (*RegisterHandlerRequest)(nil), // 4: cortex.axon.RegisterHandlerRequest - (*HandlerInvokeOption)(nil), // 5: cortex.axon.HandlerInvokeOption - (*HandlerOption)(nil), // 6: cortex.axon.HandlerOption - (*RegisterHandlerResponse)(nil), // 7: cortex.axon.RegisterHandlerResponse - (*UnregisterHandlerRequest)(nil), // 8: cortex.axon.UnregisterHandlerRequest - (*UnregisterHandlerResponse)(nil), // 9: cortex.axon.UnregisterHandlerResponse - (*ListHandlersRequest)(nil), // 10: cortex.axon.ListHandlersRequest - (*HandlerInfo)(nil), // 11: cortex.axon.HandlerInfo - (*ListHandlersResponse)(nil), // 12: cortex.axon.ListHandlersResponse - (*DispatchRequest)(nil), // 13: cortex.axon.DispatchRequest - (*DispatchMessage)(nil), // 14: cortex.axon.DispatchMessage - (*DispatchHandlerInvoke)(nil), // 15: cortex.axon.DispatchHandlerInvoke - (*InvocationFailure)(nil), // 16: cortex.axon.InvocationFailure - (*DispatchEvent)(nil), // 17: cortex.axon.DispatchEvent - (*RetryInvocationRequest)(nil), // 18: cortex.axon.RetryInvocationRequest - (*RetryInvocationResponse)(nil), // 19: cortex.axon.RetryInvocationResponse - (*Log)(nil), // 20: cortex.axon.Log - (*ReportInvocationRequest)(nil), // 21: cortex.axon.ReportInvocationRequest - (*ReportInvocationResponse)(nil), // 22: cortex.axon.ReportInvocationResponse - (*InvokeResult)(nil), // 23: cortex.axon.InvokeResult - (*GetHandlerHistoryRequest)(nil), // 24: cortex.axon.GetHandlerHistoryRequest - (*HandlerExecution)(nil), // 25: cortex.axon.HandlerExecution - (*GetHandlerHistoryResponse)(nil), // 26: cortex.axon.GetHandlerHistoryResponse - nil, // 27: cortex.axon.DispatchHandlerInvoke.ArgsEntry - (*Error)(nil), // 28: cortex.axon.Error - (*timestamppb.Timestamp)(nil), // 29: google.protobuf.Timestamp + (FailureAction)(0), // 4: cortex.axon.FailureAction + (*RegisterHandlerRequest)(nil), // 5: cortex.axon.RegisterHandlerRequest + (*HandlerInvokeOption)(nil), // 6: cortex.axon.HandlerInvokeOption + (*HandlerOption)(nil), // 7: cortex.axon.HandlerOption + (*RegisterHandlerResponse)(nil), // 8: cortex.axon.RegisterHandlerResponse + (*UnregisterHandlerRequest)(nil), // 9: cortex.axon.UnregisterHandlerRequest + (*UnregisterHandlerResponse)(nil), // 10: cortex.axon.UnregisterHandlerResponse + (*ListHandlersRequest)(nil), // 11: cortex.axon.ListHandlersRequest + (*HandlerInfo)(nil), // 12: cortex.axon.HandlerInfo + (*ListHandlersResponse)(nil), // 13: cortex.axon.ListHandlersResponse + (*DispatchRequest)(nil), // 14: cortex.axon.DispatchRequest + (*DispatchMessage)(nil), // 15: cortex.axon.DispatchMessage + (*DispatchHandlerInvoke)(nil), // 16: cortex.axon.DispatchHandlerInvoke + (*InvocationFailure)(nil), // 17: cortex.axon.InvocationFailure + (*DispatchEvent)(nil), // 18: cortex.axon.DispatchEvent + (*FailureResolution)(nil), // 19: cortex.axon.FailureResolution + (*DispatchEventsRequest)(nil), // 20: cortex.axon.DispatchEventsRequest + (*Log)(nil), // 21: cortex.axon.Log + (*ReportInvocationRequest)(nil), // 22: cortex.axon.ReportInvocationRequest + (*ReportInvocationResponse)(nil), // 23: cortex.axon.ReportInvocationResponse + (*InvokeResult)(nil), // 24: cortex.axon.InvokeResult + (*GetHandlerHistoryRequest)(nil), // 25: cortex.axon.GetHandlerHistoryRequest + (*HandlerExecution)(nil), // 26: cortex.axon.HandlerExecution + (*GetHandlerHistoryResponse)(nil), // 27: cortex.axon.GetHandlerHistoryResponse + nil, // 28: cortex.axon.DispatchHandlerInvoke.ArgsEntry + (*Error)(nil), // 29: cortex.axon.Error + (*timestamppb.Timestamp)(nil), // 30: google.protobuf.Timestamp } var file_cortex_axon_agent_proto_depIdxs = []int32{ - 6, // 0: cortex.axon.RegisterHandlerRequest.options:type_name -> cortex.axon.HandlerOption + 7, // 0: cortex.axon.RegisterHandlerRequest.options:type_name -> cortex.axon.HandlerOption 0, // 1: cortex.axon.HandlerInvokeOption.type:type_name -> cortex.axon.HandlerInvokeType - 5, // 2: cortex.axon.HandlerOption.invoke:type_name -> cortex.axon.HandlerInvokeOption - 28, // 3: cortex.axon.RegisterHandlerResponse.error:type_name -> cortex.axon.Error - 28, // 4: cortex.axon.UnregisterHandlerResponse.error:type_name -> cortex.axon.Error - 6, // 5: cortex.axon.HandlerInfo.options:type_name -> cortex.axon.HandlerOption - 29, // 6: cortex.axon.HandlerInfo.last_invoked_client_timestamp:type_name -> google.protobuf.Timestamp - 28, // 7: cortex.axon.ListHandlersResponse.error:type_name -> cortex.axon.Error - 11, // 8: cortex.axon.ListHandlersResponse.handlers:type_name -> cortex.axon.HandlerInfo + 6, // 2: cortex.axon.HandlerOption.invoke:type_name -> cortex.axon.HandlerInvokeOption + 29, // 3: cortex.axon.RegisterHandlerResponse.error:type_name -> cortex.axon.Error + 29, // 4: cortex.axon.UnregisterHandlerResponse.error:type_name -> cortex.axon.Error + 7, // 5: cortex.axon.HandlerInfo.options:type_name -> cortex.axon.HandlerOption + 30, // 6: cortex.axon.HandlerInfo.last_invoked_client_timestamp:type_name -> google.protobuf.Timestamp + 29, // 7: cortex.axon.ListHandlersResponse.error:type_name -> cortex.axon.Error + 12, // 8: cortex.axon.ListHandlersResponse.handlers:type_name -> cortex.axon.HandlerInfo 1, // 9: cortex.axon.DispatchMessage.type:type_name -> cortex.axon.DispatchMessageType - 15, // 10: cortex.axon.DispatchMessage.invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 16, // 10: cortex.axon.DispatchMessage.invoke:type_name -> cortex.axon.DispatchHandlerInvoke 0, // 11: cortex.axon.DispatchHandlerInvoke.reason:type_name -> cortex.axon.HandlerInvokeType - 27, // 12: cortex.axon.DispatchHandlerInvoke.args:type_name -> cortex.axon.DispatchHandlerInvoke.ArgsEntry - 29, // 13: cortex.axon.DispatchHandlerInvoke.deadline:type_name -> google.protobuf.Timestamp + 28, // 12: cortex.axon.DispatchHandlerInvoke.args:type_name -> cortex.axon.DispatchHandlerInvoke.ArgsEntry + 30, // 13: cortex.axon.DispatchHandlerInvoke.deadline:type_name -> google.protobuf.Timestamp 2, // 14: cortex.axon.InvocationFailure.reason:type_name -> cortex.axon.FailureReason 3, // 15: cortex.axon.InvocationFailure.source:type_name -> cortex.axon.FailureSource - 15, // 16: cortex.axon.DispatchEvent.invoke:type_name -> cortex.axon.DispatchHandlerInvoke - 16, // 17: cortex.axon.DispatchEvent.failure:type_name -> cortex.axon.InvocationFailure - 15, // 18: cortex.axon.RetryInvocationRequest.invoke:type_name -> cortex.axon.DispatchHandlerInvoke - 28, // 19: cortex.axon.RetryInvocationResponse.error:type_name -> cortex.axon.Error - 29, // 20: cortex.axon.Log.timestamp:type_name -> google.protobuf.Timestamp - 15, // 21: cortex.axon.ReportInvocationRequest.handler_invoke:type_name -> cortex.axon.DispatchHandlerInvoke - 29, // 22: cortex.axon.ReportInvocationRequest.start_client_timestamp:type_name -> google.protobuf.Timestamp - 23, // 23: cortex.axon.ReportInvocationRequest.result:type_name -> cortex.axon.InvokeResult - 28, // 24: cortex.axon.ReportInvocationRequest.error:type_name -> cortex.axon.Error - 20, // 25: cortex.axon.ReportInvocationRequest.logs:type_name -> cortex.axon.Log - 28, // 26: cortex.axon.ReportInvocationResponse.error:type_name -> cortex.axon.Error - 29, // 27: cortex.axon.GetHandlerHistoryRequest.start_time:type_name -> google.protobuf.Timestamp - 29, // 28: cortex.axon.GetHandlerHistoryRequest.end_time:type_name -> google.protobuf.Timestamp - 29, // 29: cortex.axon.HandlerExecution.publish_server_timestamp:type_name -> google.protobuf.Timestamp - 29, // 30: cortex.axon.HandlerExecution.receive_server_timestamp:type_name -> google.protobuf.Timestamp - 29, // 31: cortex.axon.HandlerExecution.start_client_timestamp:type_name -> google.protobuf.Timestamp - 28, // 32: cortex.axon.HandlerExecution.error:type_name -> cortex.axon.Error - 20, // 33: cortex.axon.HandlerExecution.logs:type_name -> cortex.axon.Log - 28, // 34: cortex.axon.GetHandlerHistoryResponse.error:type_name -> cortex.axon.Error - 25, // 35: cortex.axon.GetHandlerHistoryResponse.history:type_name -> cortex.axon.HandlerExecution - 4, // 36: cortex.axon.AxonAgent.RegisterHandler:input_type -> cortex.axon.RegisterHandlerRequest - 8, // 37: cortex.axon.AxonAgent.UnregisterHandler:input_type -> cortex.axon.UnregisterHandlerRequest - 10, // 38: cortex.axon.AxonAgent.ListHandlers:input_type -> cortex.axon.ListHandlersRequest - 24, // 39: cortex.axon.AxonAgent.GetHandlerHistory:input_type -> cortex.axon.GetHandlerHistoryRequest - 13, // 40: cortex.axon.AxonAgent.Dispatch:input_type -> cortex.axon.DispatchRequest - 21, // 41: cortex.axon.AxonAgent.ReportInvocation:input_type -> cortex.axon.ReportInvocationRequest - 13, // 42: cortex.axon.AxonAgent.DispatchEvents:input_type -> cortex.axon.DispatchRequest - 18, // 43: cortex.axon.AxonAgent.RetryInvocation:input_type -> cortex.axon.RetryInvocationRequest - 7, // 44: cortex.axon.AxonAgent.RegisterHandler:output_type -> cortex.axon.RegisterHandlerResponse - 9, // 45: cortex.axon.AxonAgent.UnregisterHandler:output_type -> cortex.axon.UnregisterHandlerResponse - 12, // 46: cortex.axon.AxonAgent.ListHandlers:output_type -> cortex.axon.ListHandlersResponse - 26, // 47: cortex.axon.AxonAgent.GetHandlerHistory:output_type -> cortex.axon.GetHandlerHistoryResponse - 14, // 48: cortex.axon.AxonAgent.Dispatch:output_type -> cortex.axon.DispatchMessage - 22, // 49: cortex.axon.AxonAgent.ReportInvocation:output_type -> cortex.axon.ReportInvocationResponse - 17, // 50: cortex.axon.AxonAgent.DispatchEvents:output_type -> cortex.axon.DispatchEvent - 19, // 51: cortex.axon.AxonAgent.RetryInvocation:output_type -> cortex.axon.RetryInvocationResponse - 44, // [44:52] is the sub-list for method output_type - 36, // [36:44] is the sub-list for method input_type - 36, // [36:36] is the sub-list for extension type_name - 36, // [36:36] is the sub-list for extension extendee - 0, // [0:36] is the sub-list for field type_name + 16, // 16: cortex.axon.DispatchEvent.invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 17, // 17: cortex.axon.DispatchEvent.failure:type_name -> cortex.axon.InvocationFailure + 4, // 18: cortex.axon.FailureResolution.action:type_name -> cortex.axon.FailureAction + 19, // 19: cortex.axon.DispatchEventsRequest.resolution:type_name -> cortex.axon.FailureResolution + 30, // 20: cortex.axon.Log.timestamp:type_name -> google.protobuf.Timestamp + 16, // 21: cortex.axon.ReportInvocationRequest.handler_invoke:type_name -> cortex.axon.DispatchHandlerInvoke + 30, // 22: cortex.axon.ReportInvocationRequest.start_client_timestamp:type_name -> google.protobuf.Timestamp + 24, // 23: cortex.axon.ReportInvocationRequest.result:type_name -> cortex.axon.InvokeResult + 29, // 24: cortex.axon.ReportInvocationRequest.error:type_name -> cortex.axon.Error + 17, // 25: cortex.axon.ReportInvocationRequest.failure:type_name -> cortex.axon.InvocationFailure + 21, // 26: cortex.axon.ReportInvocationRequest.logs:type_name -> cortex.axon.Log + 29, // 27: cortex.axon.ReportInvocationResponse.error:type_name -> cortex.axon.Error + 30, // 28: cortex.axon.GetHandlerHistoryRequest.start_time:type_name -> google.protobuf.Timestamp + 30, // 29: cortex.axon.GetHandlerHistoryRequest.end_time:type_name -> google.protobuf.Timestamp + 30, // 30: cortex.axon.HandlerExecution.publish_server_timestamp:type_name -> google.protobuf.Timestamp + 30, // 31: cortex.axon.HandlerExecution.receive_server_timestamp:type_name -> google.protobuf.Timestamp + 30, // 32: cortex.axon.HandlerExecution.start_client_timestamp:type_name -> google.protobuf.Timestamp + 29, // 33: cortex.axon.HandlerExecution.error:type_name -> cortex.axon.Error + 21, // 34: cortex.axon.HandlerExecution.logs:type_name -> cortex.axon.Log + 29, // 35: cortex.axon.GetHandlerHistoryResponse.error:type_name -> cortex.axon.Error + 26, // 36: cortex.axon.GetHandlerHistoryResponse.history:type_name -> cortex.axon.HandlerExecution + 5, // 37: cortex.axon.AxonAgent.RegisterHandler:input_type -> cortex.axon.RegisterHandlerRequest + 9, // 38: cortex.axon.AxonAgent.UnregisterHandler:input_type -> cortex.axon.UnregisterHandlerRequest + 11, // 39: cortex.axon.AxonAgent.ListHandlers:input_type -> cortex.axon.ListHandlersRequest + 25, // 40: cortex.axon.AxonAgent.GetHandlerHistory:input_type -> cortex.axon.GetHandlerHistoryRequest + 14, // 41: cortex.axon.AxonAgent.Dispatch:input_type -> cortex.axon.DispatchRequest + 22, // 42: cortex.axon.AxonAgent.ReportInvocation:input_type -> cortex.axon.ReportInvocationRequest + 20, // 43: cortex.axon.AxonAgent.DispatchEvents:input_type -> cortex.axon.DispatchEventsRequest + 8, // 44: cortex.axon.AxonAgent.RegisterHandler:output_type -> cortex.axon.RegisterHandlerResponse + 10, // 45: cortex.axon.AxonAgent.UnregisterHandler:output_type -> cortex.axon.UnregisterHandlerResponse + 13, // 46: cortex.axon.AxonAgent.ListHandlers:output_type -> cortex.axon.ListHandlersResponse + 27, // 47: cortex.axon.AxonAgent.GetHandlerHistory:output_type -> cortex.axon.GetHandlerHistoryResponse + 15, // 48: cortex.axon.AxonAgent.Dispatch:output_type -> cortex.axon.DispatchMessage + 23, // 49: cortex.axon.AxonAgent.ReportInvocation:output_type -> cortex.axon.ReportInvocationResponse + 18, // 50: cortex.axon.AxonAgent.DispatchEvents:output_type -> cortex.axon.DispatchEvent + 44, // [44:51] is the sub-list for method output_type + 37, // [37:44] is the sub-list for method input_type + 37, // [37:37] is the sub-list for extension type_name + 37, // [37:37] is the sub-list for extension extendee + 0, // [0:37] is the sub-list for field type_name } func init() { file_cortex_axon_agent_proto_init() } @@ -1980,7 +2070,7 @@ func file_cortex_axon_agent_proto_init() { File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_cortex_axon_agent_proto_rawDesc), len(file_cortex_axon_agent_proto_rawDesc)), - NumEnums: 4, + NumEnums: 5, NumMessages: 24, NumExtensions: 0, NumServices: 1, diff --git a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go index 83f80606..4125fee0 100644 --- a/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go +++ b/sdks/go/.generated/proto/github.com/cortexapps/axon/cortex-axon-agent_grpc.pb.go @@ -26,7 +26,6 @@ const ( AxonAgent_Dispatch_FullMethodName = "/cortex.axon.AxonAgent/Dispatch" AxonAgent_ReportInvocation_FullMethodName = "/cortex.axon.AxonAgent/ReportInvocation" AxonAgent_DispatchEvents_FullMethodName = "/cortex.axon.AxonAgent/DispatchEvents" - AxonAgent_RetryInvocation_FullMethodName = "/cortex.axon.AxonAgent/RetryInvocation" ) // AxonAgentClient is the client API for AxonAgent service. @@ -39,13 +38,12 @@ type AxonAgentClient interface { GetHandlerHistory(ctx context.Context, in *GetHandlerHistoryRequest, opts ...grpc.CallOption) (*GetHandlerHistoryResponse, error) Dispatch(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[DispatchRequest, DispatchMessage], error) ReportInvocation(ctx context.Context, in *ReportInvocationRequest, opts ...grpc.CallOption) (*ReportInvocationResponse, error) - // DispatchEvents streams failures the agent detected on its own: invocations - // that never reached the client, or that the client never reported. They use - // their own stream so a client that has stopped reading Dispatch because - // every handler is busy still hears about them. - DispatchEvents(ctx context.Context, in *DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[DispatchEvent], error) - // RetryInvocation queues a failed invocation again as a new attempt. - RetryInvocation(ctx context.Context, in *RetryInvocationRequest, opts ...grpc.CallOption) (*RetryInvocationResponse, error) + // DispatchEvents sends the client every failed invocation, and the client + // answers each with what to do about it. It is its own stream so a client + // that has stopped reading Dispatch because every handler is busy still + // hears about failures. An event the client has not answered when the + // stream ends is sent again on the next one. + DispatchEvents(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[DispatchEventsRequest, DispatchEvent], error) } type axonAgentClient struct { @@ -119,34 +117,18 @@ func (c *axonAgentClient) ReportInvocation(ctx context.Context, in *ReportInvoca return out, nil } -func (c *axonAgentClient) DispatchEvents(ctx context.Context, in *DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[DispatchEvent], error) { +func (c *axonAgentClient) DispatchEvents(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[DispatchEventsRequest, DispatchEvent], error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) stream, err := c.cc.NewStream(ctx, &AxonAgent_ServiceDesc.Streams[1], AxonAgent_DispatchEvents_FullMethodName, cOpts...) if err != nil { return nil, err } - x := &grpc.GenericClientStream[DispatchRequest, DispatchEvent]{ClientStream: stream} - if err := x.ClientStream.SendMsg(in); err != nil { - return nil, err - } - if err := x.ClientStream.CloseSend(); err != nil { - return nil, err - } + x := &grpc.GenericClientStream[DispatchEventsRequest, DispatchEvent]{ClientStream: stream} return x, nil } // This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type AxonAgent_DispatchEventsClient = grpc.ServerStreamingClient[DispatchEvent] - -func (c *axonAgentClient) RetryInvocation(ctx context.Context, in *RetryInvocationRequest, opts ...grpc.CallOption) (*RetryInvocationResponse, error) { - cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - out := new(RetryInvocationResponse) - err := c.cc.Invoke(ctx, AxonAgent_RetryInvocation_FullMethodName, in, out, cOpts...) - if err != nil { - return nil, err - } - return out, nil -} +type AxonAgent_DispatchEventsClient = grpc.BidiStreamingClient[DispatchEventsRequest, DispatchEvent] // AxonAgentServer is the server API for AxonAgent service. // All implementations must embed UnimplementedAxonAgentServer @@ -158,13 +140,12 @@ type AxonAgentServer interface { GetHandlerHistory(context.Context, *GetHandlerHistoryRequest) (*GetHandlerHistoryResponse, error) Dispatch(grpc.BidiStreamingServer[DispatchRequest, DispatchMessage]) error ReportInvocation(context.Context, *ReportInvocationRequest) (*ReportInvocationResponse, error) - // DispatchEvents streams failures the agent detected on its own: invocations - // that never reached the client, or that the client never reported. They use - // their own stream so a client that has stopped reading Dispatch because - // every handler is busy still hears about them. - DispatchEvents(*DispatchRequest, grpc.ServerStreamingServer[DispatchEvent]) error - // RetryInvocation queues a failed invocation again as a new attempt. - RetryInvocation(context.Context, *RetryInvocationRequest) (*RetryInvocationResponse, error) + // DispatchEvents sends the client every failed invocation, and the client + // answers each with what to do about it. It is its own stream so a client + // that has stopped reading Dispatch because every handler is busy still + // hears about failures. An event the client has not answered when the + // stream ends is sent again on the next one. + DispatchEvents(grpc.BidiStreamingServer[DispatchEventsRequest, DispatchEvent]) error mustEmbedUnimplementedAxonAgentServer() } @@ -193,12 +174,9 @@ func (UnimplementedAxonAgentServer) Dispatch(grpc.BidiStreamingServer[DispatchRe func (UnimplementedAxonAgentServer) ReportInvocation(context.Context, *ReportInvocationRequest) (*ReportInvocationResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method ReportInvocation not implemented") } -func (UnimplementedAxonAgentServer) DispatchEvents(*DispatchRequest, grpc.ServerStreamingServer[DispatchEvent]) error { +func (UnimplementedAxonAgentServer) DispatchEvents(grpc.BidiStreamingServer[DispatchEventsRequest, DispatchEvent]) error { return status.Errorf(codes.Unimplemented, "method DispatchEvents not implemented") } -func (UnimplementedAxonAgentServer) RetryInvocation(context.Context, *RetryInvocationRequest) (*RetryInvocationResponse, error) { - return nil, status.Errorf(codes.Unimplemented, "method RetryInvocation not implemented") -} func (UnimplementedAxonAgentServer) mustEmbedUnimplementedAxonAgentServer() {} func (UnimplementedAxonAgentServer) testEmbeddedByValue() {} @@ -318,33 +296,11 @@ func _AxonAgent_ReportInvocation_Handler(srv interface{}, ctx context.Context, d } func _AxonAgent_DispatchEvents_Handler(srv interface{}, stream grpc.ServerStream) error { - m := new(DispatchRequest) - if err := stream.RecvMsg(m); err != nil { - return err - } - return srv.(AxonAgentServer).DispatchEvents(m, &grpc.GenericServerStream[DispatchRequest, DispatchEvent]{ServerStream: stream}) + return srv.(AxonAgentServer).DispatchEvents(&grpc.GenericServerStream[DispatchEventsRequest, DispatchEvent]{ServerStream: stream}) } // This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type AxonAgent_DispatchEventsServer = grpc.ServerStreamingServer[DispatchEvent] - -func _AxonAgent_RetryInvocation_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { - in := new(RetryInvocationRequest) - if err := dec(in); err != nil { - return nil, err - } - if interceptor == nil { - return srv.(AxonAgentServer).RetryInvocation(ctx, in) - } - info := &grpc.UnaryServerInfo{ - Server: srv, - FullMethod: AxonAgent_RetryInvocation_FullMethodName, - } - handler := func(ctx context.Context, req interface{}) (interface{}, error) { - return srv.(AxonAgentServer).RetryInvocation(ctx, req.(*RetryInvocationRequest)) - } - return interceptor(ctx, in, info, handler) -} +type AxonAgent_DispatchEventsServer = grpc.BidiStreamingServer[DispatchEventsRequest, DispatchEvent] // AxonAgent_ServiceDesc is the grpc.ServiceDesc for AxonAgent service. // It's only intended for direct use with grpc.RegisterService, @@ -373,10 +329,6 @@ var AxonAgent_ServiceDesc = grpc.ServiceDesc{ MethodName: "ReportInvocation", Handler: _AxonAgent_ReportInvocation_Handler, }, - { - MethodName: "RetryInvocation", - Handler: _AxonAgent_RetryInvocation_Handler, - }, }, Streams: []grpc.StreamDesc{ { @@ -389,6 +341,7 @@ var AxonAgent_ServiceDesc = grpc.ServiceDesc{ StreamName: "DispatchEvents", Handler: _AxonAgent_DispatchEvents_Handler, ServerStreams: true, + ClientStreams: true, }, }, Metadata: "cortex-axon-agent.proto", diff --git a/sdks/go/README.md b/sdks/go/README.md index 4dec063c..85ad3c9b 100644 --- a/sdks/go/README.md +++ b/sdks/go/README.md @@ -114,6 +114,8 @@ This will begin executing your handler every second. `WithFailureHook` runs your own code for every failed invocation, including ones only the agent sees (such as a webhook dropped because every handler was busy). Use it to send failed messages somewhere you can replay them from. +Every failure goes through the agent: it sends each one to the hook and does +what the hook returns. This needs an agent that supports failure events. ```go agent := axon.NewAxonAgent( @@ -123,7 +125,7 @@ agent := axon.NewAxonAgent( if failure.Source == pb.FailureSource_RUNTIME && info.Attempt < 3 { return axon.FailureActionRetry } - return axon.FailureActionIgnore + return axon.FailureActionDrop }), ) ``` @@ -132,5 +134,5 @@ agent := axon.NewAxonAgent( `failure.Source` says whether handler code (`APPLICATION`) or Axon (`RUNTIME`) failed. A timeout during a Cortex API call counts as `RUNTIME`. `CallerRejected` means the sender was told the message wasn't accepted (a -webhook got a 503). `FailureActionRetry` queues the invocation again as its -next attempt, up to 5 attempts. +webhook got a 503). `FailureActionRetry` has the agent queue the invocation +again as its next attempt, up to 5 attempts; `FailureActionDrop` drops it. diff --git a/sdks/go/axon_agent.go b/sdks/go/axon_agent.go index c1c195dd..4af2279f 100644 --- a/sdks/go/axon_agent.go +++ b/sdks/go/axon_agent.go @@ -544,16 +544,15 @@ func (a *Agent) invokeHandler(ctx context.Context, invoke *pb.DispatchHandlerInv zap.Any("error", report.GetError()), ) } + // The agent sends the failure back to the failure hook, so it is + // handled the same way as failures only the agent can see. + if failure != nil { + report.Failure = failure.toProto() + } _, err := a.client.agent().ReportInvocation(context.Background(), report) if err != nil { a.logger.Error("failed to report invocation", zap.Error(err)) } - - // After the report, so a slow hook can't make the agent decide this - // invocation was never answered and report the failure a second time. - if failure != nil { - a.runFailureHook(newInvocationInfo(invoke), failure) - } } func (a *Agent) executeHandlerWithRecover(handler *handlerInfo, ctx HandlerContext) (result any, d time.Duration, err error) { diff --git a/sdks/go/hooks.go b/sdks/go/hooks.go index f3faf7ee..f2c7c90a 100644 --- a/sdks/go/hooks.go +++ b/sdks/go/hooks.go @@ -57,18 +57,37 @@ func failureFromProto(f *pb.InvocationFailure) *Failure { } } -// FailureAction is what a FailureHook wants done with a failed invocation. +func (f *Failure) toProto() *pb.InvocationFailure { + return &pb.InvocationFailure{ + Reason: f.Reason, + Source: f.Source, + Message: f.Message, + CallerRejected: f.CallerRejected, + } +} + +// FailureAction is what a FailureHook wants the agent to do with a failed +// invocation. type FailureAction int const ( - // FailureActionIgnore leaves the failure as it is. - FailureActionIgnore FailureAction = iota - // FailureActionRetry queues the invocation again as its next attempt. + // FailureActionDrop does nothing more with the invocation. + FailureActionDrop FailureAction = iota + // FailureActionRetry has the agent queue the invocation again as its + // next attempt. FailureActionRetry ) -// FailureHook runs for every failed invocation, including failures the agent -// detects on its own, such as an invocation that never reached this process. +func (a FailureAction) toProto() pb.FailureAction { + if a == FailureActionRetry { + return pb.FailureAction_FAILURE_ACTION_RETRY + } + return pb.FailureAction_FAILURE_ACTION_DROP +} + +// FailureHook runs for every failed invocation, wherever it failed: in this +// process, or in the agent, such as an invocation that never reached this +// process. The agent sends each failure here and applies the answer. type FailureHook func(ctx context.Context, info InvocationInfo, failure *Failure) FailureAction // WithFailureHook sets a hook that runs for every failed invocation. @@ -78,55 +97,17 @@ func WithFailureHook(hook FailureHook) Option { } } -func (a *Agent) runFailureHook(info InvocationInfo, failure *Failure) { - if a.onFailure == nil { - return - } - if a.callFailureHook(info, failure) == FailureActionRetry { - a.retryInvocation(info) - } -} - -func (a *Agent) callFailureHook(info InvocationInfo, failure *Failure) (action FailureAction) { +func (a *Agent) runFailureHook(info InvocationInfo, failure *Failure) (action FailureAction) { defer func() { if r := recover(); r != nil { - a.logger.Error("panic in failure hook", zap.Any("panic", r)) - action = FailureActionIgnore + a.logger.Error("panic in failure hook, dropping the invocation", + zap.String("invocation-id", info.InvocationID), zap.Any("panic", r)) + action = FailureActionDrop } }() return a.onFailure(context.Background(), info, failure) } -func (a *Agent) retryInvocation(info InvocationInfo) { - stub := a.client.agent() - if stub == nil { - a.logger.Error("cannot retry invocation, no agent connection", - zap.String("invocation-id", info.InvocationID)) - return - } - resp, err := stub.RetryInvocation(context.Background(), &pb.RetryInvocationRequest{Invoke: info.invoke}) - if status.Code(err) == codes.Unimplemented { - a.logger.Warn("agent does not support retries, upgrade the agent", - zap.String("invocation-id", info.InvocationID)) - return - } - if err != nil { - a.logger.Error("failed to retry invocation", - zap.String("invocation-id", info.InvocationID), zap.Error(err)) - return - } - if resp.GetError() != nil { - a.logger.Warn("agent rejected retry", - zap.String("invocation-id", info.InvocationID), - zap.String("error", resp.GetError().GetMessage())) - return - } - a.logger.Info("retrying invocation", - zap.String("invocation-id", info.InvocationID), - zap.String("retry-invocation-id", resp.GetInvocationId()), - zap.Int("attempt", info.Attempt+1)) -} - var errEventsUnsupported = errors.New("agent does not support failure events") // runFailureEvents reads the agent's failure events until ctx is done. It has @@ -135,7 +116,7 @@ func (a *Agent) runFailureEvents(ctx context.Context) { for ctx.Err() == nil { err := a.readFailureEvents(ctx) if errors.Is(err, errEventsUnsupported) { - a.logger.Warn("agent does not send failure events, upgrade the agent to receive them") + a.logger.Warn("agent does not send failure events, upgrade the agent to use a failure hook") return } if ctx.Err() != nil { @@ -149,15 +130,20 @@ func (a *Agent) runFailureEvents(ctx context.Context) { } } +// readFailureEvents runs the hook for each failure the agent sends and +// answers with what it returned. func (a *Agent) readFailureEvents(ctx context.Context) error { stub := a.client.agent() if stub == nil { return fmt.Errorf("failed to create agent connection") } - stream, err := stub.DispatchEvents(ctx, &pb.DispatchRequest{ + stream, err := stub.DispatchEvents(ctx) + if err != nil { + return err + } + err = stream.Send(&pb.DispatchEventsRequest{ DispatchId: a.DispatchId, ClientVersion: version.Client, - FailureEvents: true, }) if err != nil { return err @@ -173,6 +159,15 @@ func (a *Agent) readFailureEvents(ctx context.Context) error { if ev.GetInvoke() == nil { continue } - a.runFailureHook(newInvocationInfo(ev.GetInvoke()), failureFromProto(ev.GetFailure())) + action := a.runFailureHook(newInvocationInfo(ev.GetInvoke()), failureFromProto(ev.GetFailure())) + err = stream.Send(&pb.DispatchEventsRequest{ + Resolution: &pb.FailureResolution{ + InvocationId: ev.GetInvoke().GetInvocationId(), + Action: action.toProto(), + }, + }) + if err != nil { + return err + } } } diff --git a/sdks/go/hooks_test.go b/sdks/go/hooks_test.go index 0b3d1548..02655e82 100644 --- a/sdks/go/hooks_test.go +++ b/sdks/go/hooks_test.go @@ -3,6 +3,7 @@ package axon import ( "context" "fmt" + "io" "sync" "testing" "time" @@ -16,186 +17,146 @@ import ( "google.golang.org/protobuf/types/known/timestamppb" ) -type hookCall struct { - info InvocationInfo - failure *Failure -} - -// hookRecorder is a failure hook that records each call. -type hookRecorder struct { - mu sync.Mutex - calls []hookCall - action FailureAction - seen chan struct{} -} - -func newHookRecorder(action FailureAction) *hookRecorder { - return &hookRecorder{action: action, seen: make(chan struct{}, 10)} -} - -func (r *hookRecorder) hook(_ context.Context, info InvocationInfo, failure *Failure) FailureAction { - r.mu.Lock() - r.calls = append(r.calls, hookCall{info, failure}) - r.mu.Unlock() - r.seen <- struct{}{} - return r.action -} - -func (r *hookRecorder) only(t *testing.T) hookCall { - r.mu.Lock() - defer r.mu.Unlock() - require.Len(t, r.calls, 1) - return r.calls[0] -} - -// eventStream is a DispatchEvents stream that delivers events and then -// blocks until the agent stops reading. -type eventStream struct { - grpc.ClientStream - ctx context.Context - events chan *pb.DispatchEvent - err error -} - -func (s *eventStream) Recv() (*pb.DispatchEvent, error) { - if s.err != nil { - return nil, s.err - } - select { - case ev := <-s.events: - return ev, nil - case <-s.ctx.Done(): - return nil, s.ctx.Err() - } -} - -func expectEventStream(mock *mockGrpcClient, events ...*pb.DispatchEvent) { - mock.agentStub.EXPECT().DispatchEvents(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn( - func(ctx context.Context, req *pb.DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[pb.DispatchEvent], error) { - ch := make(chan *pb.DispatchEvent, len(events)) - for _, ev := range events { - ch <- ev - } - events = nil - return &eventStream{ctx: ctx, events: ch}, nil - }) -} - -func runWithFailureHook(t *testing.T, handler any, timeoutMs int, rec *hookRecorder, setup ...func(*mockGrpcClient)) { +// reportedFailure runs handler once and returns the failure it reported to +// the agent, or nil. +func reportedFailure(t *testing.T, handler any, timeoutMs int, setup ...func(*mockGrpcClient)) *pb.InvocationFailure { controller := gomock.NewController(t) defer controller.Finish() - executeHandlerWithAgent(t, controller, handler, timeoutMs, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.onFailure = rec.hook - expectEventStream(mock) - for _, s := range setup { - s(mock) - } - }) + var reported *pb.ReportInvocationRequest + executeHandlerWithAgent(t, controller, handler, timeoutMs, + func(req *pb.ReportInvocationRequest) { reported = req }, + func(_ *Agent, mock *mockGrpcClient) { + for _, s := range setup { + s(mock) + } + }) + require.NotNil(t, reported) + return reported.GetFailure() } -func TestFailureHookNotCalledOnSuccess(t *testing.T) { - rec := newHookRecorder(FailureActionRetry) - runWithFailureHook(t, func(ctx HandlerContext) error { return nil }, 0, rec) - - rec.mu.Lock() - defer rec.mu.Unlock() - require.Empty(t, rec.calls) +func TestSuccessReportsNoFailure(t *testing.T) { + require.Nil(t, reportedFailure(t, func(ctx HandlerContext) error { return nil }, 0)) } -func TestFailureHookOnHandlerError(t *testing.T) { - rec := newHookRecorder(FailureActionIgnore) - runWithFailureHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec) - - call := rec.only(t) - require.Equal(t, "func1", call.info.HandlerName) - require.Equal(t, pb.HandlerInvokeType_RUN_NOW, call.info.InvokeType) - require.Equal(t, 1, call.info.Attempt, "an agent that doesn't send attempt means the first") - failure := call.failure +func TestHandlerErrorIsReportedAsApplicationFailure(t *testing.T) { + failure := reportedFailure(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0) require.NotNil(t, failure) require.Equal(t, pb.FailureReason_HANDLER_ERROR, failure.Reason) require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) require.Equal(t, "ohnoes", failure.Message) } -func TestFailureHookOnPanic(t *testing.T) { - rec := newHookRecorder(FailureActionIgnore) - runWithFailureHook(t, func(ctx HandlerContext) error { panic("ohnoes") }, 0, rec) - - failure := rec.only(t).failure +func TestPanicIsReportedAsHandlerError(t *testing.T) { + failure := reportedFailure(t, func(ctx HandlerContext) error { panic("ohnoes") }, 0) require.Equal(t, pb.FailureReason_HANDLER_ERROR, failure.Reason) require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) } -func TestFailureHookOnHandlerTimeout(t *testing.T) { - rec := newHookRecorder(FailureActionIgnore) - runWithFailureHook(t, func(ctx HandlerContext) error { +func TestSlowHandlerIsReportedAsApplicationTimeout(t *testing.T) { + failure := reportedFailure(t, func(ctx HandlerContext) error { time.Sleep(50 * time.Millisecond) return nil - }, 1, rec) - - failure := rec.only(t).failure + }, 1) require.Equal(t, pb.FailureReason_TIMEOUT, failure.Reason) require.Equal(t, pb.FailureSource_APPLICATION, failure.Source) } -func TestFailureHookTimeoutDuringApiCallIsRuntime(t *testing.T) { - rec := newHookRecorder(FailureActionIgnore) - runWithFailureHook(t, func(ctx HandlerContext) error { +func TestTimeoutDuringApiCallIsReportedAsRuntime(t *testing.T) { + failure := reportedFailure(t, func(ctx HandlerContext) error { _, err := ctx.CortexJsonApiCall("GET", "/api/v1/slow", "") return err - }, 20, rec, func(mock *mockGrpcClient) { + }, 20, func(mock *mockGrpcClient) { mock.apiStub.EXPECT().Call(gomock.Any(), gomock.Any()).DoAndReturn( func(ctx context.Context, req *pb.CallRequest, opts ...grpc.CallOption) (*pb.CallResponse, error) { <-ctx.Done() return nil, status.FromContextError(ctx.Err()).Err() }) }) - - failure := rec.only(t).failure require.Equal(t, pb.FailureReason_TIMEOUT, failure.Reason) require.Equal(t, pb.FailureSource_RUNTIME, failure.Source) } -func TestFailureHookRetry(t *testing.T) { - rec := newHookRecorder(FailureActionRetry) - retried := make(chan *pb.RetryInvocationRequest, 1) - runWithFailureHook(t, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, rec, func(mock *mockGrpcClient) { - mock.agentStub.EXPECT().RetryInvocation(gomock.Any(), gomock.Any()).DoAndReturn( - func(ctx context.Context, req *pb.RetryInvocationRequest, opts ...grpc.CallOption) (*pb.RetryInvocationResponse, error) { - retried <- req - return &pb.RetryInvocationResponse{InvocationId: "retry-1"}, nil - }) - }) - - req := <-retried - require.Equal(t, "func1", req.Invoke.HandlerName) -} - -func TestFailureHookPanicIsContained(t *testing.T) { +func TestFailureHookNotRunLocally(t *testing.T) { controller := gomock.NewController(t) defer controller.Finish() - - reported := false - executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, - func(*pb.ReportInvocationRequest) { reported = true }, + called := false + executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { return fmt.Errorf("ohnoes") }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { panic("hook bug") } + // The agent sends the failure back as an event; it is not run here. + agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { + called = true + return FailureActionDrop + } expectEventStream(mock) }) - require.True(t, reported) + require.False(t, called) +} + +// eventStream is a DispatchEvents stream: it delivers events, records the +// client's messages, and then blocks until the client stops reading. +type eventStream struct { + grpc.ClientStream + ctx context.Context + events chan *pb.DispatchEvent + err error + + mu sync.Mutex + sent []*pb.DispatchEventsRequest + // answered gets each resolution the client sends. + answered chan *pb.FailureResolution } -func TestFailureEventReachesFailureHook(t *testing.T) { - rec := newHookRecorder(FailureActionRetry) - retried := make(chan *pb.RetryInvocationRequest, 1) +func (s *eventStream) Send(req *pb.DispatchEventsRequest) error { + s.mu.Lock() + s.sent = append(s.sent, req) + s.mu.Unlock() + if req.Resolution != nil { + s.answered <- req.Resolution + } + return nil +} - ev := &pb.DispatchEvent{ +func (s *eventStream) Recv() (*pb.DispatchEvent, error) { + if s.err != nil { + return nil, s.err + } + select { + case ev := <-s.events: + return ev, nil + case <-s.ctx.Done(): + return nil, io.EOF + } +} + +func (s *eventStream) CloseSend() error { return nil } + +// expectEventStream serves events on the first DispatchEvents stream the +// client opens, and returns the stream so the test can see its answers. +func expectEventStream(mock *mockGrpcClient, events ...*pb.DispatchEvent) *eventStream { + stream := &eventStream{ + ctx: context.Background(), + events: make(chan *pb.DispatchEvent, len(events)), + answered: make(chan *pb.FailureResolution, len(events)+1), + } + for _, ev := range events { + stream.events <- ev + } + mock.agentStub.EXPECT().DispatchEvents(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn( + func(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[pb.DispatchEventsRequest, pb.DispatchEvent], error) { + stream.ctx = ctx + return stream, nil + }) + return stream +} + +func failureEvent(id string, attempt int32) *pb.DispatchEvent { + return &pb.DispatchEvent{ Invoke: &pb.DispatchHandlerInvoke{ - InvocationId: "lost-1", + InvocationId: id, HandlerName: "func1", Reason: pb.HandlerInvokeType_WEBHOOK, - Args: map[string]string{"body": "payload"}, - Attempt: 2, + Args: map[string]string{"body": "payload-" + id}, + Attempt: attempt, }, Failure: &pb.InvocationFailure{ Reason: pb.FailureReason_NO_RESOURCES, @@ -204,32 +165,84 @@ func TestFailureEventReachesFailureHook(t *testing.T) { CallerRejected: true, }, } +} +// runWithEvents runs one handler invocation while the agent sends events. +// The handler holds the run open until every event has been answered. +func runWithEvents(t *testing.T, hook FailureHook, events ...*pb.DispatchEvent) *eventStream { controller := gomock.NewController(t) defer controller.Finish() + var stream *eventStream + answeredAll := make(chan struct{}) executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { - // Hold the run open until the event has been handled. - <-rec.seen + select { + case <-answeredAll: + case <-time.After(5 * time.Second): + } return nil }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.onFailure = rec.hook - expectEventStream(mock, ev) - mock.agentStub.EXPECT().RetryInvocation(gomock.Any(), gomock.Any()).DoAndReturn( - func(ctx context.Context, req *pb.RetryInvocationRequest, opts ...grpc.CallOption) (*pb.RetryInvocationResponse, error) { - retried <- req - return &pb.RetryInvocationResponse{InvocationId: "retry-1"}, nil - }) + agent.onFailure = hook + stream = expectEventStream(mock, events...) + go func() { + for { + stream.mu.Lock() + n := 0 + for _, req := range stream.sent { + if req.Resolution != nil { + n++ + } + } + stream.mu.Unlock() + if n >= len(events) { + close(answeredAll) + return + } + time.Sleep(5 * time.Millisecond) + } + }() }) + return stream +} + +func TestFailureEventRunsHookAndAnswers(t *testing.T) { + var got InvocationInfo + var gotFailure *Failure + stream := runWithEvents(t, func(_ context.Context, info InvocationInfo, failure *Failure) FailureAction { + got, gotFailure = info, failure + return FailureActionRetry + }, failureEvent("lost-1", 2)) + + require.Equal(t, "lost-1", got.InvocationID) + require.Equal(t, "payload-lost-1", got.Args["body"]) + require.Equal(t, 2, got.Attempt) + require.Equal(t, pb.FailureReason_NO_RESOURCES, gotFailure.Reason) + require.True(t, gotFailure.CallerRejected) + + stream.mu.Lock() + defer stream.mu.Unlock() + require.Equal(t, "", stream.sent[0].GetResolution().GetInvocationId(), "the first message only identifies the client") + require.NotEmpty(t, stream.sent[0].DispatchId) + res := stream.sent[1].GetResolution() + require.Equal(t, "lost-1", res.InvocationId) + require.Equal(t, pb.FailureAction_FAILURE_ACTION_RETRY, res.Action) +} + +func TestFailureHookDropIsAnswered(t *testing.T) { + stream := runWithEvents(t, func(context.Context, InvocationInfo, *Failure) FailureAction { + return FailureActionDrop + }, failureEvent("lost-1", 1)) + require.Equal(t, pb.FailureAction_FAILURE_ACTION_DROP, (<-stream.answered).Action) +} - rec.mu.Lock() - first := rec.calls[0] - rec.mu.Unlock() - require.Equal(t, "lost-1", first.info.InvocationID) - require.Equal(t, "payload", first.info.Args["body"]) - require.Equal(t, 2, first.info.Attempt) - require.Equal(t, pb.FailureReason_NO_RESOURCES, first.failure.Reason) - require.True(t, first.failure.CallerRejected) - require.Equal(t, "lost-1", (<-retried).Invoke.InvocationId) +func TestFailureHookPanicDrops(t *testing.T) { + stream := runWithEvents(t, func(context.Context, InvocationInfo, *Failure) FailureAction { + panic("hook bug") + }, failureEvent("lost-1", 1), failureEvent("lost-2", 1)) + // The panic is contained, and the stream keeps going. + require.Equal(t, "lost-1", (<-stream.answered).InvocationId) + second := <-stream.answered + require.Equal(t, "lost-2", second.InvocationId) + require.Equal(t, pb.FailureAction_FAILURE_ACTION_DROP, second.Action) } func TestFailureEventsStopOnOlderAgent(t *testing.T) { @@ -240,10 +253,10 @@ func TestFailureEventsStopOnOlderAgent(t *testing.T) { time.Sleep(50 * time.Millisecond) return nil }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { - agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { return FailureActionIgnore } + agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { return FailureActionDrop } // Times(1): an agent without the RPC must not be redialed in a loop. mock.agentStub.EXPECT().DispatchEvents(gomock.Any(), gomock.Any()).Times(1).Return( - &eventStream{err: status.Error(codes.Unimplemented, "unknown method DispatchEvents")}, nil) + &eventStream{err: status.Error(codes.Unimplemented, "unknown method DispatchEvents"), answered: make(chan *pb.FailureResolution, 1)}, nil) }) } @@ -257,7 +270,7 @@ func TestDispatchRequestAdvertisesFailureEvents(t *testing.T) { executeHandlerWithAgent(t, controller, func(ctx HandlerContext) error { return nil }, 0, nil, func(agent *Agent, mock *mockGrpcClient) { if withHook { - agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { return FailureActionIgnore } + agent.onFailure = func(context.Context, InvocationInfo, *Failure) FailureAction { return FailureActionDrop } expectEventStream(mock) } sent = mock.lastStream diff --git a/sdks/go/mock_axon/mocks.go b/sdks/go/mock_axon/mocks.go index 40fd1196..738818ec 100644 --- a/sdks/go/mock_axon/mocks.go +++ b/sdks/go/mock_axon/mocks.go @@ -107,22 +107,22 @@ func (mr *MockAxonAgentClientMockRecorder) Dispatch(ctx any, opts ...any) *gomoc } // DispatchEvents mocks base method. -func (m *MockAxonAgentClient) DispatchEvents(ctx context.Context, in *axon.DispatchRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[axon.DispatchEvent], error) { +func (m *MockAxonAgentClient) DispatchEvents(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[axon.DispatchEventsRequest, axon.DispatchEvent], error) { m.ctrl.T.Helper() - varargs := []any{ctx, in} + varargs := []any{ctx} for _, a := range opts { varargs = append(varargs, a) } ret := m.ctrl.Call(m, "DispatchEvents", varargs...) - ret0, _ := ret[0].(grpc.ServerStreamingClient[axon.DispatchEvent]) + ret0, _ := ret[0].(grpc.BidiStreamingClient[axon.DispatchEventsRequest, axon.DispatchEvent]) ret1, _ := ret[1].(error) return ret0, ret1 } // DispatchEvents indicates an expected call of DispatchEvents. -func (mr *MockAxonAgentClientMockRecorder) DispatchEvents(ctx, in any, opts ...any) *gomock.Call { +func (mr *MockAxonAgentClientMockRecorder) DispatchEvents(ctx any, opts ...any) *gomock.Call { mr.mock.ctrl.T.Helper() - varargs := append([]any{ctx, in}, opts...) + varargs := append([]any{ctx}, opts...) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchEvents", reflect.TypeOf((*MockAxonAgentClient)(nil).DispatchEvents), varargs...) } @@ -206,26 +206,6 @@ func (mr *MockAxonAgentClientMockRecorder) ReportInvocation(ctx, in any, opts .. return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReportInvocation", reflect.TypeOf((*MockAxonAgentClient)(nil).ReportInvocation), varargs...) } -// RetryInvocation mocks base method. -func (m *MockAxonAgentClient) RetryInvocation(ctx context.Context, in *axon.RetryInvocationRequest, opts ...grpc.CallOption) (*axon.RetryInvocationResponse, error) { - m.ctrl.T.Helper() - varargs := []any{ctx, in} - for _, a := range opts { - varargs = append(varargs, a) - } - ret := m.ctrl.Call(m, "RetryInvocation", varargs...) - ret0, _ := ret[0].(*axon.RetryInvocationResponse) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// RetryInvocation indicates an expected call of RetryInvocation. -func (mr *MockAxonAgentClientMockRecorder) RetryInvocation(ctx, in any, opts ...any) *gomock.Call { - mr.mock.ctrl.T.Helper() - varargs := append([]any{ctx, in}, opts...) - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RetryInvocation", reflect.TypeOf((*MockAxonAgentClient)(nil).RetryInvocation), varargs...) -} - // UnregisterHandler mocks base method. func (m *MockAxonAgentClient) UnregisterHandler(ctx context.Context, in *axon.UnregisterHandlerRequest, opts ...grpc.CallOption) (*axon.UnregisterHandlerResponse, error) { m.ctrl.T.Helper() diff --git a/sdks/python/README.md b/sdks/python/README.md index a332d7b2..b42c503d 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -111,10 +111,13 @@ Set it to `1` for the old behaviour, where handlers run one at a time. Pass `on_failure` to `AxonClient` to run your own code whenever an invocation fails, for example to send the failed message somewhere you can -replay it from. It also runs for failures that only the agent can see, such -as a webhook the agent couldn't hand to this process because every handler -was busy. Those failures arrive on their own stream, so they get through even -while every handler slot is taken. +replay it from. Every failure goes through the agent, wherever it happened: +a handler that raised in this process, or a webhook the agent couldn't hand +to this process because every handler was busy. The agent sends each one to +your hook and does what the hook returns. Failures arrive on their own +stream, so they get through even while every handler slot is taken, and a +failure your hook hadn't answered when the connection dropped is sent again. +This needs an agent that supports failure events. ```python from cortex_axon import AxonClient, FailureAction, FailureSource @@ -155,9 +158,10 @@ client.run() `caller_rejected` is true when the sender was told the message wasn't accepted (a webhook got a 503), so the sender may also retry it. -Returning `FailureAction.RETRY` asks the agent to queue the invocation -again, with `attempt` incremented. The agent stops after 5 attempts. -Returning nothing ignores the failure. +Returning `FailureAction.RETRY` has the agent queue the invocation again, +with `attempt` incremented; it stops after 5 attempts. Returning +`FailureAction.DROP`, or nothing, drops it. If the hook raises, the +invocation is dropped. The hook runs for every handler type, not only webhooks. Python can't stop a handler that overruns its timeout: at the deadline the SDK reports the diff --git a/sdks/python/cortex_axon/axon_client.py b/sdks/python/cortex_axon/axon_client.py index 536a18cc..34804fbb 100644 --- a/sdks/python/cortex_axon/axon_client.py +++ b/sdks/python/cortex_axon/axon_client.py @@ -273,36 +273,24 @@ def _find_handler(self, name): return h return None - def _run_failure_hook(self, info: InvocationInfo, failure: Failure, stub): - if not self.on_failure: - return + def _run_failure_hook(self, info: InvocationInfo, failure: Failure) -> FailureAction: try: action = self.on_failure(info, failure) except Exception: - print(f"Error in on_failure hook for {info.handler_name} (id={info.invocation_id})") + print(f"Error in on_failure hook for {info.handler_name} (id={info.invocation_id}), dropping it") traceback.print_exc() - return - if action == FailureAction.RETRY: - self._retry_invocation(info, stub) - - def _retry_invocation(self, info: InvocationInfo, stub): - try: - response = stub.RetryInvocation( - cortex_axon_agent_pb2.RetryInvocationRequest(invoke=info.invoke) + return FailureAction.DROP + return action or FailureAction.DROP + + def _answer_failure(self, event, answers: queue.Queue): + info = InvocationInfo.from_proto(event.invoke) + action = self._run_failure_hook(info, Failure.from_proto(event.failure)) + answers.put(cortex_axon_agent_pb2.DispatchEventsRequest( + resolution=cortex_axon_agent_pb2.FailureResolution( + invocation_id=info.invocation_id, + action=action.to_proto(), ) - except grpc.RpcError as e: - if e.code() == grpc.StatusCode.UNIMPLEMENTED: - print("Agent does not support retries, upgrade the agent") - else: - print(f"Error retrying invocation {info.invocation_id}: {e}") - return - if response.error and response.error.code: - print(f"Agent rejected retry of {info.invocation_id}: {response.error.message}") - return - print( - f"Retrying {info.handler_name} (id={info.invocation_id}) " - f"as id={response.invocation_id} attempt={info.attempt + 1}" - ) + )) def _start_failure_events(self): if not self.on_failure or self._events_thread: @@ -327,30 +315,32 @@ def _run_failure_events(self): stub = cortex_axon_agent_pb2_grpc.AxonAgentStub(channel) try: while not self._events_stop.is_set(): - try: - self._events_call = stub.DispatchEvents( - cortex_axon_agent_pb2.DispatchRequest( - dispatch_id=self.id, - client_version=AGENT_VERSION, - failure_events=True, - ) + # Answers for this stream only: one the agent never got is + # sent again on the next stream, and answered there. + answers = queue.Queue() + + def requests(answers=answers): + yield cortex_axon_agent_pb2.DispatchEventsRequest( + dispatch_id=self.id, client_version=AGENT_VERSION ) + while (answer := answers.get()) is not None: + yield answer + + try: + self._events_call = stub.DispatchEvents(requests()) for event in self._events_call: if not event.HasField("invoke"): continue - self._hook_executor.submit( - self._run_failure_hook, - InvocationInfo.from_proto(event.invoke), - Failure.from_proto(event.failure), - stub, - ) + self._hook_executor.submit(self._answer_failure, event, answers) except grpc.RpcError as e: if self._events_stop.is_set(): return if e.code() == grpc.StatusCode.UNIMPLEMENTED: - print("Agent does not send failure events, upgrade the agent to receive them") + print("Agent does not send failure events, upgrade the agent to use on_failure") return print(f"Failure event stream ended ({e.code()}), reconnecting in {RECONNECT_WAIT_SECONDS} seconds...") + finally: + answers.put(None) self._events_stop.wait(RECONNECT_WAIT_SECONDS) finally: channel.close() @@ -391,7 +381,6 @@ def _invoke_handler(self, handler_to_invoke, invoke, stub): handler_result = None failure = None result = None - info = InvocationInfo.from_proto(invoke) now = datetime.now() timeout, capped, _ = _invoke_timeout(invoke, time()) ctx = self._handler_context(invoke.args, timeout * 1000 if timeout else 0) @@ -428,11 +417,9 @@ def on_timeout(): start_client_timestamp=now, duration_ms=invoke.timeout_ms, error=common_pb2.Error(code="timeout", message=message), + failure=Failure(reason=FailureReason.TIMEOUT, source=source, message=message).to_proto(), logs=list(ctx.logs), )) - self._run_failure_hook( - info, Failure(reason=FailureReason.TIMEOUT, source=source, message=message), stub - ) timer = None if timeout: @@ -481,16 +468,14 @@ def on_timeout(): duration_ms=duration_ms, result=handler_result, error=handler_err, + # The agent sends the failure back to on_failure, so it is + # handled the same way as failures only the agent can see. + failure=failure.to_proto() if failure else None, logs=ctx.logs, ) self._report(stub, invoke_info) - # After the report, so a slow hook can't make the agent decide this - # invocation was never answered and report the failure a second time. - if failure is not None: - self._run_failure_hook(info, failure, stub) - def _report(self, stub, invoke_info): # A failure here must not escape the worker: the dispatch loop is no # longer in the call path, so an exception would be swallowed by the diff --git a/sdks/python/cortex_axon/hooks.py b/sdks/python/cortex_axon/hooks.py index f4240d8c..d9d57cb7 100644 --- a/sdks/python/cortex_axon/hooks.py +++ b/sdks/python/cortex_axon/hooks.py @@ -25,11 +25,16 @@ class FailureSource(IntEnum): class FailureAction(Enum): - # Leave the failure as it is. - IGNORE = "ignore" - # Queue the invocation again as its next attempt. + # Do nothing more with the invocation. + DROP = "drop" + # Have the agent queue the invocation again as its next attempt. RETRY = "retry" + def to_proto(self) -> int: + if self == FailureAction.RETRY: + return cortex_axon_agent_pb2.FAILURE_ACTION_RETRY + return cortex_axon_agent_pb2.FAILURE_ACTION_DROP + @dataclass(frozen=True) class InvocationInfo: @@ -65,6 +70,14 @@ class Failure: # sender got a 503), so it may arrive again. caller_rejected: bool = False + def to_proto(self) -> cortex_axon_agent_pb2.InvocationFailure: + return cortex_axon_agent_pb2.InvocationFailure( + reason=self.reason, + source=self.source, + message=self.message, + caller_rejected=self.caller_rejected, + ) + @classmethod def from_proto(cls, failure: cortex_axon_agent_pb2.InvocationFailure) -> "Failure": return cls( @@ -75,7 +88,8 @@ def from_proto(cls, failure: cortex_axon_agent_pb2.InvocationFailure) -> "Failur ) -# Runs for every failed invocation, including failures the agent detects on -# its own, such as an invocation that never reached this process. Returning -# None means FailureAction.IGNORE. +# Runs for every failed invocation, wherever it failed: in this process, or in +# the agent, such as an invocation that never reached this process. The agent +# sends each failure here and applies the answer. Returning None means +# FailureAction.DROP. FailureHook = Callable[[InvocationInfo, Failure], Optional[FailureAction]] diff --git a/sdks/python/example/main.py b/sdks/python/example/main.py index f4ded6a4..e03c87a0 100644 --- a/sdks/python/example/main.py +++ b/sdks/python/example/main.py @@ -60,7 +60,7 @@ def on_failure(info, failure): # A real hook would save info.args somewhere it can be replayed from. if info.attempt < 3: return FailureAction.RETRY - return FailureAction.IGNORE + return FailureAction.DROP def run(): diff --git a/sdks/python/tests/test_hooks.py b/sdks/python/tests/test_hooks.py index 1f23b524..759612b1 100644 --- a/sdks/python/tests/test_hooks.py +++ b/sdks/python/tests/test_hooks.py @@ -33,7 +33,8 @@ def __init__(self, invokes=(), events=(), events_code=None): self.events = list(events) self.events_code = events_code self.dispatch_requests = queue.Queue() - self.retries = queue.Queue() + self.events_requests = queue.Queue() + self.resolutions = queue.Queue() self.events_calls = 0 self.reports = queue.Queue() self.stop = threading.Event() @@ -55,10 +56,20 @@ def Dispatch(self, request_iterator, context): self.invokes = [] self.stop.wait(timeout=10) - def DispatchEvents(self, request, context): + def DispatchEvents(self, request_iterator, context): self.events_calls += 1 if self.events_code: context.abort(self.events_code, "not here") + self.events_requests.put(next(request_iterator)) + + def read_answers(): + # Ends with an error when the client goes away at teardown. + with contextlib.suppress(Exception): + for request in request_iterator: + if request.HasField("resolution"): + self.resolutions.put(request.resolution) + + threading.Thread(target=read_answers, daemon=True).start() events, self.events = self.events, [] yield from events self.stop.wait(timeout=10) @@ -67,10 +78,6 @@ def ReportInvocation(self, request, context): self.reports.put(request) return cortex_axon_agent_pb2.ReportInvocationResponse() - def RetryInvocation(self, request, context): - self.retries.put(request) - return cortex_axon_agent_pb2.RetryInvocationResponse(invocation_id="retry-1") - class SlowApi(cortex_api_pb2_grpc.CortexApiServicer): def __init__(self): @@ -95,6 +102,20 @@ def invoke(invocation_id="1", timeout_ms=0, deadline_in=None, **args): return msg +def failure_event(invocation_id, attempt=1, **args): + msg = invoke(invocation_id, **args) + msg.attempt = attempt + return cortex_axon_agent_pb2.DispatchEvent( + invoke=msg, + failure=cortex_axon_agent_pb2.InvocationFailure( + reason=cortex_axon_agent_pb2.NO_RESOURCES, + source=cortex_axon_agent_pb2.RUNTIME, + message="dispatch queue is full", + caller_rejected=True, + ), + ) + + def serve(servicer, add): server = grpc.server(ThreadPoolExecutor(max_workers=8)) port = server.add_insecure_port("127.0.0.1:0") @@ -159,22 +180,23 @@ def start(agent, handler, hooks=None, api=None, max_concurrent_handlers=4): client._events_thread.join(timeout=5) -def test_failure_hook_not_called_on_success(run_client): - hooks = Hooks(action=FailureAction.RETRY) +def reported_failure(agent): + report = agent.reports.get(timeout=5) + return report.failure if report.HasField("failure") else None + + +def test_success_reports_no_failure(run_client): agent = HookAgent(invokes=[invoke(body="payload")]) def my_hooked_webhook(ctx): return "ok" - run_client(agent, my_hooked_webhook, hooks) + run_client(agent, my_hooked_webhook, Hooks()) - assert agent.reports.get(timeout=5).result.value == "ok" - with pytest.raises(queue.Empty): - hooks.failures.get(timeout=0.5) - assert agent.retries.empty() + assert reported_failure(agent) is None -def test_failure_hook_on_handler_error_and_retry(run_client): +def test_handler_error_is_reported_not_hooked_locally(run_client): hooks = Hooks(action=FailureAction.RETRY) agent = HookAgent(invokes=[invoke(body="payload")]) @@ -183,21 +205,16 @@ def my_hooked_webhook(ctx): run_client(agent, my_hooked_webhook, hooks) - info, failure = hooks.failures.get(timeout=5) - assert info.invocation_id == "1" - assert info.args == {"body": "payload"} - assert info.attempt == 1 + failure = reported_failure(agent) assert failure.reason == FailureReason.HANDLER_ERROR assert failure.source == FailureSource.APPLICATION assert failure.message == "ohnoes" - - retry = agent.retries.get(timeout=5) - assert retry.invoke.invocation_id == "1" - assert retry.invoke.args["body"] == "payload" + # The agent sends it back as an event; the SDK doesn't run the hook itself. + with pytest.raises(queue.Empty): + hooks.failures.get(timeout=0.5) -def test_api_timeout_is_runtime(run_client): - hooks = Hooks() +def test_api_timeout_is_reported_as_runtime(run_client): api = SlowApi() agent = HookAgent(invokes=[invoke(timeout_ms=200)]) @@ -205,16 +222,15 @@ def my_hooked_webhook(ctx): ctx.cortex_api_call("/api/v1/slow", "GET") try: - run_client(agent, my_hooked_webhook, hooks, api=api) - _, failure = hooks.failures.get(timeout=5) + run_client(agent, my_hooked_webhook, Hooks(), api=api) + failure = reported_failure(agent) finally: api.release.set() assert failure.reason == FailureReason.TIMEOUT assert failure.source == FailureSource.RUNTIME -def test_slow_handler_times_out_as_application(run_client): - hooks = Hooks() +def test_slow_handler_is_reported_as_application_timeout(run_client): agent = HookAgent(invokes=[invoke(timeout_ms=200)]) finished = threading.Event() @@ -222,25 +238,22 @@ def my_hooked_webhook(ctx): time.sleep(0.6) finished.set() - run_client(agent, my_hooked_webhook, hooks) + run_client(agent, my_hooked_webhook, Hooks()) - _, failure = hooks.failures.get(timeout=5) - assert failure.reason == FailureReason.TIMEOUT - assert failure.source == FailureSource.APPLICATION report = agent.reports.get(timeout=5) assert report.error.code == "timeout" + assert report.failure.reason == FailureReason.TIMEOUT + assert report.failure.source == FailureSource.APPLICATION - # The late finish is dropped: no second report, no second hook call. + # The late finish is dropped: no second report. assert finished.wait(timeout=5) - with pytest.raises(queue.Empty): - hooks.failures.get(timeout=0.5) + time.sleep(0.2) assert agent.reports.empty() def test_expired_invocation_is_not_started(run_client): """The agent has already reported an invocation past its deadline, so the SDK must not run it: a failure hook may be replaying it.""" - hooks = Hooks() ran = queue.Queue() agent = HookAgent(invokes=[ invoke("stale", timeout_ms=10000, deadline_in=0.5), @@ -250,50 +263,70 @@ def test_expired_invocation_is_not_started(run_client): def my_hooked_webhook(ctx): ran.put(ctx.args) - run_client(agent, my_hooked_webhook, hooks) + run_client(agent, my_hooked_webhook, Hooks()) assert agent.reports.get(timeout=5).handler_invoke.invocation_id == "fresh" - # The stale one neither ran nor reached the hook: the agent reports it. - with pytest.raises(queue.Empty): - hooks.failures.get(timeout=0.5) + time.sleep(0.5) assert agent.reports.empty() assert ran.qsize() == 1 def test_timeout_capped_by_agent_deadline_is_runtime(run_client): - hooks = Hooks() # Its own timeout is long, but the agent gives up in 1.5s: it waited. agent = HookAgent(invokes=[invoke(timeout_ms=60000, deadline_in=1.5)]) def my_hooked_webhook(ctx): time.sleep(2) - run_client(agent, my_hooked_webhook, hooks) + run_client(agent, my_hooked_webhook, Hooks()) - _, failure = hooks.failures.get(timeout=5) + failure = reported_failure(agent) assert failure.reason == FailureReason.TIMEOUT assert failure.source == FailureSource.RUNTIME - assert agent.reports.get(timeout=5).error.code == "timeout" -def test_hook_errors_are_contained(run_client): - agent = HookAgent(invokes=[invoke("1"), invoke("2")]) - ran = queue.Queue() +def test_failure_event_runs_hook_and_answers(run_client): + hooks = Hooks(action=FailureAction.RETRY) + agent = HookAgent(events=[failure_event("lost-1", attempt=2, body="payload")]) + + run_client(agent, lambda ctx: None, hooks) + + first = agent.events_requests.get(timeout=5) + assert first.dispatch_id + assert not first.HasField("resolution") + + info, failure = hooks.failures.get(timeout=5) + assert info.invocation_id == "lost-1" + assert info.args == {"body": "payload"} + assert info.attempt == 2 + assert failure.reason == FailureReason.NO_RESOURCES + assert failure.source == FailureSource.RUNTIME + assert failure.caller_rejected + + answer = agent.resolutions.get(timeout=5) + assert answer.invocation_id == "lost-1" + assert answer.action == cortex_axon_agent_pb2.FAILURE_ACTION_RETRY + + +def test_hook_returning_nothing_drops(run_client): + agent = HookAgent(events=[failure_event("lost-1")]) + run_client(agent, lambda ctx: None, Hooks(action=None)) + + assert agent.resolutions.get(timeout=5).action == cortex_axon_agent_pb2.FAILURE_ACTION_DROP + + +def test_hook_errors_drop_and_are_contained(run_client): + agent = HookAgent(events=[failure_event("lost-1"), failure_event("lost-2")]) class Broken(Hooks): def on_failure(self, info, failure): raise RuntimeError("hook bug") - def my_hooked_webhook(ctx): - ran.put(True) - raise ValueError("ohnoes") + run_client(agent, lambda ctx: None, Broken()) - run_client(agent, my_hooked_webhook, Broken()) - # Both still run and report, despite the hook raising each time. - ran.get(timeout=5) - ran.get(timeout=5) - assert agent.reports.get(timeout=5).error.code == "unexpected" - assert agent.reports.get(timeout=5).error.code == "unexpected" + answers = [agent.resolutions.get(timeout=5) for _ in range(2)] + assert {a.invocation_id for a in answers} == {"lost-1", "lost-2"} + assert all(a.action == cortex_axon_agent_pb2.FAILURE_ACTION_DROP for a in answers) def test_dispatch_request_advertises_failure_events(run_client): @@ -313,18 +346,7 @@ def test_failure_events_reach_hook_while_handlers_are_saturated(run_client): slots = 2 release = threading.Event() started = threading.Barrier(slots + 1) - events = [ - cortex_axon_agent_pb2.DispatchEvent( - invoke=invoke(f"lost-{i}", body=f"msg-{i}"), - failure=cortex_axon_agent_pb2.InvocationFailure( - reason=cortex_axon_agent_pb2.NO_RESOURCES, - source=cortex_axon_agent_pb2.RUNTIME, - message="dispatch queue is full", - caller_rejected=True, - ), - ) - for i in range(20) - ] + events = [failure_event(f"lost-{i}", body=f"msg-{i}") for i in range(20)] # More invocations than slots, so the dispatch reader blocks. agent = HookAgent(invokes=[invoke(str(i)) for i in range(slots + 2)], events=events) hooks = Hooks(action=FailureAction.RETRY) @@ -341,12 +363,8 @@ def my_hooked_webhook(ctx): seen = [hooks.failures.get(timeout=5) for _ in events] bodies = sorted(info.args["body"] for info, _ in seen) assert bodies == sorted(f"msg-{i}" for i in range(20)) - for _, failure in seen: - assert failure.reason == FailureReason.NO_RESOURCES - assert failure.source == FailureSource.RUNTIME - assert failure.caller_rejected - retried = {agent.retries.get(timeout=5).invoke.invocation_id for _ in events} - assert retried == {f"lost-{i}" for i in range(20)} + answered = {agent.resolutions.get(timeout=5).invocation_id for _ in events} + assert answered == {f"lost-{i}" for i in range(20)} finally: release.set() started.abort()