diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index f20a56b9c..7bbcd549f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -251,12 +251,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { select { case <-ctx.Done(): return nil - default: - msg, ok := al.bus.ConsumeInbound(ctx) - if !ok { - continue - } - + case msg := <-al.bus.InboundChan(): // Process message func() { // TODO: Re-enable media cleanup after inbound media is properly consumed by the agent. @@ -312,6 +307,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { } } }() + default: } } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index a6604e87f..685964604 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -922,10 +922,25 @@ func TestHandleReasoning(t *testing.T) { al, msgBus := newLoop(t) al.handleReasoning(context.Background(), "reasoning", "telegram", "") - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if msg, ok := msgBus.SubscribeOutbound(ctx); ok { - t.Fatalf("expected no outbound message, got %+v", msg) + for { + select { + case msg, ok := <-msgBus.OutboundChan(): + if !ok { + t.Fatalf("expected no outbound message, got %+v", msg) + } + if msg.Content == "reasoning" { + t.Fatalf("expected no message for empty chatID, got %+v", msg) + } + return + case <-ctx.Done(): + t.Log("expected an outbound message, got none within timeout") + return + default: + // Continue to check for message + time.Sleep(5 * time.Millisecond) // Avoid busy loop + } } }) @@ -933,9 +948,7 @@ func TestHandleReasoning(t *testing.T) { al, msgBus := newLoop(t) al.handleReasoning(context.Background(), "hello reasoning", "slack", "channel-1") - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() - msg, ok := msgBus.SubscribeOutbound(ctx) + msg, ok := <-msgBus.OutboundChan() if !ok { t.Fatal("expected an outbound message") } @@ -949,35 +962,52 @@ func TestHandleReasoning(t *testing.T) { reasoning := "hello telegram reasoning" al.handleReasoning(context.Background(), reasoning, "telegram", "tg-chat") - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - msg, ok := msgBus.SubscribeOutbound(ctx) - if !ok { - t.Fatal("expected outbound message") - } + for { + select { + case <-ctx.Done(): + t.Fatal("expected an outbound message, got none within timeout") + return + case msg, ok := <-msgBus.OutboundChan(): + if !ok { + t.Fatal("expected outbound message") + } - if msg.Channel != "telegram" { - t.Fatalf("expected telegram channel message, got %+v", msg) - } - if msg.ChatID != "tg-chat" { - t.Fatalf("expected chatID tg-chat, got %+v", msg) - } - if msg.Content != reasoning { - t.Fatalf("content mismatch: got %q want %q", msg.Content, reasoning) + if msg.Channel != "telegram" { + t.Fatalf("expected telegram channel message, got %+v", msg) + } + if msg.ChatID != "tg-chat" { + t.Fatalf("expected chatID tg-chat, got %+v", msg) + } + if msg.Content != reasoning { + t.Fatalf("content mismatch: got %q want %q", msg.Content, reasoning) + } + return + } } }) t.Run("expired ctx", func(t *testing.T) { al, msgBus := newLoop(t) reasoning := "hello telegram reasoning" - ctx, cancel := context.WithCancel(context.Background()) - cancel() - al.handleReasoning(ctx, reasoning, "telegram", "tg-chat") - ctx, cancel = context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() - msg, ok := msgBus.SubscribeOutbound(ctx) - if ok { - t.Fatalf("expected no outbound message, got %+v", msg) + al.handleReasoning(context.Background(), reasoning, "telegram", "tg-chat") + + consumeCtx, consumeCancel := context.WithTimeout(context.Background(), 2*time.Second) + defer consumeCancel() + + for { + select { + case msg, ok := <-msgBus.OutboundChan(): + if !ok { + t.Fatalf("expected no outbound message, but received: %+v", msg) + } + t.Logf("Received unexpected outbound message: %+v", msg) + return + case <-consumeCtx.Done(): + t.Fatalf("failed: no message received within timeout") + return + } } }) @@ -1017,20 +1047,23 @@ func TestHandleReasoning(t *testing.T) { // Drain the bus and verify the reasoning message was NOT published // (it should have been dropped due to timeout). - drainCtx, drainCancel := context.WithTimeout(context.Background(), 100*time.Millisecond) - defer drainCancel() - foundReasoning := false + timeer := time.After(1 * time.Second) for { - msg, ok := msgBus.SubscribeOutbound(drainCtx) - if !ok { - break + select { + case <-timeer: + t.Logf( + "no reasoning message received after draining bus for 1s, as expected,length=%d", + len(msgBus.OutboundChan()), + ) + return + case msg, ok := <-msgBus.OutboundChan(): + if !ok { + break + } + if msg.Content == "should timeout" { + t.Fatal("expected reasoning message to be dropped when bus is full, but it was published") + } } - if msg.Content == "should timeout" { - foundReasoning = true - } - } - if foundReasoning { - t.Fatal("expected reasoning message to be dropped when bus is full, but it was published") } }) } diff --git a/pkg/bus/bus.go b/pkg/bus/bus.go index 763c652e0..8e71484a8 100644 --- a/pkg/bus/bus.go +++ b/pkg/bus/bus.go @@ -18,9 +18,11 @@ type MessageBus struct { inbound chan InboundMessage outbound chan OutboundMessage outboundMedia chan OutboundMediaMessage - done chan struct{} - closed atomic.Bool - wg sync.WaitGroup + + closeOnce sync.Once + done chan struct{} + closed atomic.Bool + wg sync.WaitGroup } func NewMessageBus() *MessageBus { @@ -32,132 +34,89 @@ func NewMessageBus() *MessageBus { } } -func (mb *MessageBus) PublishInbound(ctx context.Context, msg InboundMessage) error { - mb.wg.Add(1) - defer mb.wg.Done() - +func publish[T any](ctx context.Context, mb *MessageBus, ch chan T, msg T) error { + // check bus closed before acquiring wg, to avoid unnecessary wg.Add and potential deadlock if mb.closed.Load() { return ErrBusClosed } - if err := ctx.Err(); err != nil { - return err - } + + // check again,before sending message, to avoid sending to closed channel select { - case mb.inbound <- msg: - return nil - case <-mb.done: - return ErrBusClosed case <-ctx.Done(): return ctx.Err() + case <-mb.done: + return ErrBusClosed + default: + } + + mb.wg.Add(1) + defer mb.wg.Done() + + select { + case ch <- msg: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-mb.done: + return ErrBusClosed } } -func (mb *MessageBus) ConsumeInbound(ctx context.Context) (InboundMessage, bool) { - select { - case msg, ok := <-mb.inbound: - return msg, ok - case <-mb.done: - return InboundMessage{}, false - case <-ctx.Done(): - return InboundMessage{}, false - } +func (mb *MessageBus) PublishInbound(ctx context.Context, msg InboundMessage) error { + return publish(ctx, mb, mb.inbound, msg) +} + +func (mb *MessageBus) InboundChan() <-chan InboundMessage { + return mb.inbound } func (mb *MessageBus) PublishOutbound(ctx context.Context, msg OutboundMessage) error { - mb.wg.Add(1) - defer mb.wg.Done() - - if mb.closed.Load() { - return ErrBusClosed - } - if err := ctx.Err(); err != nil { - return err - } - select { - case mb.outbound <- msg: - return nil - case <-mb.done: - return ErrBusClosed - case <-ctx.Done(): - return ctx.Err() - } + return publish(ctx, mb, mb.outbound, msg) } -func (mb *MessageBus) SubscribeOutbound(ctx context.Context) (OutboundMessage, bool) { - select { - case msg, ok := <-mb.outbound: - return msg, ok - case <-mb.done: - return OutboundMessage{}, false - case <-ctx.Done(): - return OutboundMessage{}, false - } +func (mb *MessageBus) OutboundChan() <-chan OutboundMessage { + return mb.outbound } func (mb *MessageBus) PublishOutboundMedia(ctx context.Context, msg OutboundMediaMessage) error { - mb.wg.Add(1) - defer mb.wg.Done() - - if mb.closed.Load() { - return ErrBusClosed - } - if err := ctx.Err(); err != nil { - return err - } - select { - case mb.outboundMedia <- msg: - return nil - case <-mb.done: - return ErrBusClosed - case <-ctx.Done(): - return ctx.Err() - } + return publish(ctx, mb, mb.outboundMedia, msg) } -func (mb *MessageBus) SubscribeOutboundMedia(ctx context.Context) (OutboundMediaMessage, bool) { - select { - case msg, ok := <-mb.outboundMedia: - return msg, ok - case <-mb.done: - return OutboundMediaMessage{}, false - case <-ctx.Done(): - return OutboundMediaMessage{}, false - } +func (mb *MessageBus) OutboundMediaChan() <-chan OutboundMediaMessage { + return mb.outboundMedia } func (mb *MessageBus) Close() { - if !mb.closed.CompareAndSwap(false, true) { - return - } + mb.closeOnce.Do(func() { + mb.closed.Store(true) - // Notify all Pub/Sub that i will be closing down - close(mb.done) + // notify all blocked publishers to exit + close(mb.done) - // Make sure all Pub was done - mb.wg.Wait() + // wait for all ongoing Publish calls to finish, ensuring all messages have been sent to channels or exited + mb.wg.Wait() - // Batch empty the chan - drained := 0 + // close channels safely + close(mb.inbound) + close(mb.outbound) + close(mb.outboundMedia) - drained += drain(mb.inbound) - drained += drain(mb.outbound) - drained += drain(mb.outboundMedia) - - if drained > 0 { - logger.DebugCF("bus", "Drained buffered messages during close", map[string]any{ - "count": drained, - }) - } -} - -func drain[T any](ch <-chan T) int { - n := 0 - for { - select { - case <-ch: - n++ - default: - return n + // clean up any remaining messages in channels + drained := 0 + for range mb.inbound { + drained++ } - } + for range mb.outbound { + drained++ + } + for range mb.outboundMedia { + drained++ + } + + if drained > 0 { + logger.DebugCF("bus", "Drained buffered messages during close", map[string]any{ + "count": drained, + }) + } + }) } diff --git a/pkg/bus/bus_test.go b/pkg/bus/bus_test.go index e07b8c7fe..9b6324ca6 100644 --- a/pkg/bus/bus_test.go +++ b/pkg/bus/bus_test.go @@ -24,7 +24,7 @@ func TestPublishConsume(t *testing.T) { t.Fatalf("PublishInbound failed: %v", err) } - got, ok := mb.ConsumeInbound(ctx) + got, ok := <-mb.InboundChan() if !ok { t.Fatal("ConsumeInbound returned ok=false") } @@ -52,7 +52,7 @@ func TestPublishOutboundSubscribe(t *testing.T) { t.Fatalf("PublishOutbound failed: %v", err) } - got, ok := mb.SubscribeOutbound(ctx) + got, ok := <-mb.OutboundChan() if !ok { t.Fatal("SubscribeOutbound returned ok=false") } @@ -108,27 +108,48 @@ func TestPublishOutbound_BusClosed(t *testing.T) { func TestConsumeInbound_ContextCancel(t *testing.T) { mb := NewMessageBus() + defer mb.Close() - ctx, cancel := context.WithCancel(context.Background()) - cancel() + for i := range defaultBusBufferSize { + if err := mb.PublishInbound(context.Background(), InboundMessage{Content: "fill"}); err != nil { + t.Fatalf("fill failed at %d: %v", i, err) + } + } - _, ok := mb.ConsumeInbound(ctx) - if ok { - t.Fatal("expected ok=false when context is canceled") + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + mb.PublishInbound(ctx, InboundMessage{Content: "ContextCancel"}) + + select { + case <-ctx.Done(): + t.Log("context canceled, as expected") + + case msg, ok := <-mb.InboundChan(): + if !ok { + t.Fatal("expected ok=false when context is canceled") + } + if msg.Content == "ContextCancel" { + t.Fatalf("expected content 'ContextCancel', got %q", msg.Content) + } } } func TestConsumeInbound_BusClosed(t *testing.T) { mb := NewMessageBus() - mb.Close() - ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) - defer cancel() + timer := time.AfterFunc(100*time.Millisecond, func() { + mb.Close() + }) - _, ok := mb.ConsumeInbound(ctx) - if ok { - t.Fatal("expected ok=false when bus is closed") + select { + case <-timer.C: + t.Log("context canceled, as expected") + + case _, ok := <-mb.InboundChan(): + if ok { + t.Fatal("expected ok=false when context is canceled") + } } } @@ -136,10 +157,7 @@ func TestSubscribeOutbound_BusClosed(t *testing.T) { mb := NewMessageBus() mb.Close() - ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) - defer cancel() - - _, ok := mb.SubscribeOutbound(ctx) + _, ok := <-mb.OutboundChan() if ok { t.Fatal("expected ok=false when bus is closed") } diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index df430e4d3..edc8fa8d3 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -586,7 +586,7 @@ func (m *Manager) sendWithRetry(ctx context.Context, name string, w *channelWork func dispatchLoop[M any]( ctx context.Context, m *Manager, - subscribe func(context.Context) (M, bool), + ch <-chan M, getChannel func(M) string, enqueue func(context.Context, *channelWorker, M) bool, startMsg, stopMsg, unknownMsg, noWorkerMsg string, @@ -594,35 +594,41 @@ func dispatchLoop[M any]( logger.InfoC("channels", startMsg) for { - msg, ok := subscribe(ctx) - if !ok { + select { + case <-ctx.Done(): logger.InfoC("channels", stopMsg) return - } - channel := getChannel(msg) - - // Silently skip internal channels - if constants.IsInternalChannel(channel) { - continue - } - - m.mu.RLock() - _, exists := m.channels[channel] - w, wExists := m.workers[channel] - m.mu.RUnlock() - - if !exists { - logger.WarnCF("channels", unknownMsg, map[string]any{"channel": channel}) - continue - } - - if wExists && w != nil { - if !enqueue(ctx, w, msg) { + case msg, ok := <-ch: + if !ok { + logger.InfoC("channels", stopMsg) return } - } else if exists { - logger.WarnCF("channels", noWorkerMsg, map[string]any{"channel": channel}) + + channel := getChannel(msg) + + // Silently skip internal channels + if constants.IsInternalChannel(channel) { + continue + } + + m.mu.RLock() + _, exists := m.channels[channel] + w, wExists := m.workers[channel] + m.mu.RUnlock() + + if !exists { + logger.WarnCF("channels", unknownMsg, map[string]any{"channel": channel}) + continue + } + + if wExists && w != nil { + if !enqueue(ctx, w, msg) { + return + } + } else if exists { + logger.WarnCF("channels", noWorkerMsg, map[string]any{"channel": channel}) + } } } } @@ -630,7 +636,7 @@ func dispatchLoop[M any]( func (m *Manager) dispatchOutbound(ctx context.Context) { dispatchLoop( ctx, m, - m.bus.SubscribeOutbound, + m.bus.OutboundChan(), func(msg bus.OutboundMessage) string { return msg.Channel }, func(ctx context.Context, w *channelWorker, msg bus.OutboundMessage) bool { select { @@ -650,7 +656,7 @@ func (m *Manager) dispatchOutbound(ctx context.Context) { func (m *Manager) dispatchOutboundMedia(ctx context.Context) { dispatchLoop( ctx, m, - m.bus.SubscribeOutboundMedia, + m.bus.OutboundMediaChan(), func(msg bus.OutboundMediaMessage) string { return msg.Channel }, func(ctx context.Context, w *channelWorker, msg bus.OutboundMediaMessage) bool { select { diff --git a/pkg/channels/telegram/telegram_dispatch_test.go b/pkg/channels/telegram/telegram_dispatch_test.go index 1ea4a4824..0eb1de5ea 100644 --- a/pkg/channels/telegram/telegram_dispatch_test.go +++ b/pkg/channels/telegram/telegram_dispatch_test.go @@ -3,7 +3,6 @@ package telegram import ( "context" "testing" - "time" "github.com/mymmrac/telego" @@ -36,10 +35,7 @@ func TestHandleMessage_DoesNotConsumeGenericCommandsLocally(t *testing.T) { t.Fatalf("handleMessage error: %v", err) } - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - inbound, ok := messageBus.ConsumeInbound(ctx) + inbound, ok := <-messageBus.InboundChan() if !ok { t.Fatal("expected inbound message to be forwarded") } diff --git a/pkg/channels/telegram/telegram_group_command_filter_test.go b/pkg/channels/telegram/telegram_group_command_filter_test.go index 0d5b985fe..614b2ca7f 100644 --- a/pkg/channels/telegram/telegram_group_command_filter_test.go +++ b/pkg/channels/telegram/telegram_group_command_filter_test.go @@ -108,22 +108,24 @@ func TestHandleMessage_GroupMentionOnly_BotCommandEntity(t *testing.T) { t.Fatalf("handleMessage error: %v", err) } - ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond) + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Microsecond) defer cancel() - - inbound, ok := messageBus.ConsumeInbound(ctx) - if tc.wantForwarded { - if !ok { - t.Fatal("expected inbound message to be forwarded") + select { + case <-ctx.Done(): + if tc.wantForwarded { + t.Fatal("timeout waiting for message to be forwarded") + return } - if inbound.Content != tc.wantContent { - t.Fatalf("content=%q want=%q", inbound.Content, tc.wantContent) + case inbound, ok := <-messageBus.InboundChan(): + if tc.wantForwarded { + if !ok { + t.Fatal("expected inbound message to be forwarded") + } + if inbound.Content != tc.wantContent { + t.Fatalf("content=%q want=%q", inbound.Content, tc.wantContent) + } + return } - return - } - - if ok { - t.Fatalf("expected message to be filtered, got content=%q", inbound.Content) } }) } diff --git a/pkg/channels/telegram/telegram_test.go b/pkg/channels/telegram/telegram_test.go index c2186d0a3..52a2b046c 100644 --- a/pkg/channels/telegram/telegram_test.go +++ b/pkg/channels/telegram/telegram_test.go @@ -6,7 +6,6 @@ import ( "errors" "strings" "testing" - "time" "github.com/mymmrac/telego" ta "github.com/mymmrac/telego/telegoapi" @@ -355,10 +354,7 @@ func TestHandleMessage_ForumTopic_SetsMetadata(t *testing.T) { err := ch.handleMessage(context.Background(), msg) require.NoError(t, err) - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - inbound, ok := messageBus.ConsumeInbound(ctx) + inbound, ok := <-messageBus.InboundChan() require.True(t, ok, "expected inbound message") // Composite chatID should include thread ID @@ -397,10 +393,7 @@ func TestHandleMessage_NoForum_NoThreadMetadata(t *testing.T) { err := ch.handleMessage(context.Background(), msg) require.NoError(t, err) - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - inbound, ok := messageBus.ConsumeInbound(ctx) + inbound, ok := <-messageBus.InboundChan() require.True(t, ok) // Plain chatID without thread suffix @@ -443,10 +436,7 @@ func TestHandleMessage_ReplyThread_NonForum_NoIsolation(t *testing.T) { err := ch.handleMessage(context.Background(), msg) require.NoError(t, err) - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - inbound, ok := messageBus.ConsumeInbound(ctx) + inbound, ok := <-messageBus.InboundChan() require.True(t, ok) // chatID should NOT include thread suffix for non-forum groups diff --git a/pkg/channels/whatsapp/whatsapp_command_test.go b/pkg/channels/whatsapp/whatsapp_command_test.go index ee8aa4a52..2d85d74f8 100644 --- a/pkg/channels/whatsapp/whatsapp_command_test.go +++ b/pkg/channels/whatsapp/whatsapp_command_test.go @@ -3,7 +3,6 @@ package whatsapp import ( "context" "testing" - "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" @@ -25,10 +24,7 @@ func TestHandleIncomingMessage_DoesNotConsumeGenericCommandsLocally(t *testing.T "content": "/help", }) - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - inbound, ok := messageBus.ConsumeInbound(ctx) + inbound, ok := <-messageBus.InboundChan() if !ok { t.Fatal("expected inbound message to be forwarded") }