pch/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(bus): add backpressure drop budget for audio streams only

Browse source
This commit is contained in:
SiYue-ZO 2026-06-11 13:55:56 +08:00
parent 3a68d26837
commit a7208983d0
5 changed files with 414 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

225
pkg/bus/bus.go
View file

@ -3,8 +3,10 @@ package bus
import (
"context"
"errors"
"fmt"
"sync"
"sync/atomic"
"time"
runtimeevents "github.com/sipeed/picoclaw/pkg/events"
"github.com/sipeed/picoclaw/pkg/logger"
@ -13,6 +15,10 @@ import (
// ErrBusClosed is returned when publishing to a closed MessageBus.
var ErrBusClosed = errors.New("message bus closed")
// ErrBusBackpressure is returned when a publish attempt exceeds the configured
// backpressure wait budget and the message is dropped.
var ErrBusBackpressure = errors.New("message bus backpressure")
var (
ErrMissingInboundContext = errors.New("inbound message context is required")
ErrMissingOutboundContext = errors.New("outbound message context is required")
@ -21,6 +27,41 @@ var (
const defaultBusBufferSize = 64
const (
defaultAudioPublishTimeout = 150 * time.Millisecond
)
type publishPolicy struct {
stream string
// timeout is the backpressure drop budget. When positive, a full channel
// causes the message to be dropped after this duration. When zero, the
// publish blocks until context cancellation or bus close (no drop).
timeout time.Duration
}
type streamStats struct {
dropped atomic.Uint64
lastDropped atomic.Int64
lastWaitNanos atomic.Int64
}
type MessageBusStats struct {
Inbound StreamStats `json:"inbound"`
Outbound StreamStats `json:"outbound"`
OutboundMedia StreamStats `json:"outbound_media"`
AudioChunks StreamStats `json:"audio_chunks"`
VoiceControls StreamStats `json:"voice_controls"`
}
type StreamStats struct {
Depth int `json:"depth"`
Capacity int `json:"capacity"`
DroppedTotal uint64 `json:"dropped_total"`
LastDroppedAt time.Time `json:"last_dropped_at,omitempty"`
LastDropWait string `json:"last_drop_wait,omitempty"`
LastDropWaitMillis int64 `json:"last_drop_wait_ms,omitempty"`
}
// StreamDelegate is implemented by the channel Manager to provide streaming
// capabilities to the agent loop without tight coupling.
type StreamDelegate interface {
@ -62,8 +103,14 @@ type MessageBus struct {
done chan struct{}
closed atomic.Bool
wg sync.WaitGroup
publishMu sync.Mutex
streamDelegate atomic.Value // stores StreamDelegate
eventPublisher atomic.Value // stores EventPublisher
inboundStats streamStats
outboundStats streamStats
mediaStats streamStats
audioStats streamStats
voiceStats streamStats
}
// EventPublisher is the minimal runtime event publisher used by MessageBus.
@ -83,13 +130,13 @@ func NewMessageBus() *MessageBus {
}
}
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
func (mb *MessageBus) enterPublish(ctx context.Context) error {
mb.publishMu.Lock()
defer mb.publishMu.Unlock()
if mb.closed.Load() {
return ErrBusClosed
}
// check again,before sending message, to avoid sending to closed channel
select {
case <-ctx.Done():
return ctx.Err()
@ -99,8 +146,63 @@ func publish[T any](ctx context.Context, mb *MessageBus, ch chan T, msg T) error
}
mb.wg.Add(1)
return nil
}
func publish[T any](
ctx context.Context,
mb *MessageBus,
ch chan T,
msg T,
policy publishPolicy,
stats *streamStats,
scope runtimeevents.Scope,
) error {
if ctx == nil {
ctx = context.Background()
}
if err := mb.enterPublish(ctx); err != nil {
return err
}
defer mb.wg.Done()
// timeout == 0 means no backpressure drop budget; block until context
// cancellation or bus close. This is the default for critical streams
// (inbound, outbound, outboundMedia, voiceControl) where dropping
// messages silently is undesirable.
if policy.timeout > 0 {
timer := time.NewTimer(policy.timeout)
defer timer.Stop()
select {
case ch <- msg:
return nil
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
droppedTotal := stats.dropped.Add(1)
now := time.Now()
stats.lastDropped.Store(now.UnixNano())
stats.lastWaitNanos.Store(policy.timeout.Nanoseconds())
queueDepth := len(ch)
queueCap := cap(ch)
mb.publishDrop(
policy.stream, scope, "queue_full_timeout",
policy.timeout, queueDepth, queueCap, droppedTotal,
)
logger.WarnCF("bus", "Dropped bus message due to backpressure", map[string]any{
"stream": policy.stream,
"wait_ms": policy.timeout.Milliseconds(),
"queue_depth": queueDepth,
"queue_capacity": queueCap,
"dropped_total": droppedTotal,
})
return fmt.Errorf("%w: %s queue full after %s", ErrBusBackpressure, policy.stream, policy.timeout)
case <-mb.done:
return ErrBusClosed
}
}
select {
case ch <- msg:
return nil
@ -117,8 +219,13 @@ func (mb *MessageBus) PublishInbound(ctx context.Context, msg InboundMessage) er
mb.publishFailure("inbound", runtimeScopeFromInboundContext(msg.Context), ErrMissingInboundContext)
return ErrMissingInboundContext
}
if err := publish(ctx, mb, mb.inbound, msg); err != nil {
mb.publishFailure("inbound", runtimeScopeFromInboundContext(msg.Context), err)
if err := publish(ctx, mb, mb.inbound, msg, publishPolicy{
stream: "inbound",
}, &mb.inboundStats, runtimeScopeFromInboundContext(msg.Context)); err != nil {
scope := runtimeScopeFromInboundContext(msg.Context)
if !errors.Is(err, ErrBusBackpressure) {
mb.publishFailure("inbound", scope, err)
}
return err
}
return nil
@ -134,8 +241,13 @@ func (mb *MessageBus) PublishOutbound(ctx context.Context, msg OutboundMessage)
mb.publishFailure("outbound", runtimeScopeFromInboundContext(msg.Context), ErrMissingOutboundContext)
return ErrMissingOutboundContext
}
if err := publish(ctx, mb, mb.outbound, msg); err != nil {
mb.publishFailure("outbound", runtimeScopeFromInboundContext(msg.Context), err)
if err := publish(ctx, mb, mb.outbound, msg, publishPolicy{
stream: "outbound",
}, &mb.outboundStats, runtimeScopeFromInboundContext(msg.Context)); err != nil {
scope := runtimeScopeFromInboundContext(msg.Context)
if !errors.Is(err, ErrBusBackpressure) {
mb.publishFailure("outbound", scope, err)
}
return err
}
return nil
@ -151,8 +263,13 @@ func (mb *MessageBus) PublishOutboundMedia(ctx context.Context, msg OutboundMedi
mb.publishFailure("outbound_media", runtimeScopeFromInboundContext(msg.Context), ErrMissingOutboundMediaContext)
return ErrMissingOutboundMediaContext
}
if err := publish(ctx, mb, mb.outboundMedia, msg); err != nil {
mb.publishFailure("outbound_media", runtimeScopeFromInboundContext(msg.Context), err)
if err := publish(ctx, mb, mb.outboundMedia, msg, publishPolicy{
stream: "outbound_media",
}, &mb.mediaStats, runtimeScopeFromInboundContext(msg.Context)); err != nil {
scope := runtimeScopeFromInboundContext(msg.Context)
if !errors.Is(err, ErrBusBackpressure) {
mb.publishFailure("outbound_media", scope, err)
}
return err
}
return nil
@ -163,8 +280,14 @@ func (mb *MessageBus) OutboundMediaChan() <-chan OutboundMediaMessage {
}
func (mb *MessageBus) PublishAudioChunk(ctx context.Context, chunk AudioChunk) error {
if err := publish(ctx, mb, mb.audioChunks, chunk); err != nil {
mb.publishFailure("audio_chunk", runtimeScopeFromAudioChunk(chunk), err)
if err := publish(ctx, mb, mb.audioChunks, chunk, publishPolicy{
stream: "audio_chunk",
timeout: defaultAudioPublishTimeout,
}, &mb.audioStats, runtimeScopeFromAudioChunk(chunk)); err != nil {
scope := runtimeScopeFromAudioChunk(chunk)
if !errors.Is(err, ErrBusBackpressure) {
mb.publishFailure("audio_chunk", scope, err)
}
return err
}
return nil
@ -175,8 +298,13 @@ func (mb *MessageBus) AudioChunksChan() <-chan AudioChunk {
}
func (mb *MessageBus) PublishVoiceControl(ctx context.Context, ctrl VoiceControl) error {
if err := publish(ctx, mb, mb.voiceControls, ctrl); err != nil {
mb.publishFailure("voice_control", runtimeScopeFromVoiceControl(ctrl), err)
if err := publish(ctx, mb, mb.voiceControls, ctrl, publishPolicy{
stream: "voice_control",
}, &mb.voiceStats, runtimeScopeFromVoiceControl(ctrl)); err != nil {
scope := runtimeScopeFromVoiceControl(ctrl)
if !errors.Is(err, ErrBusBackpressure) {
mb.publishFailure("voice_control", scope, err)
}
return err
}
return nil
@ -204,15 +332,76 @@ func (mb *MessageBus) GetStreamer(ctx context.Context, channel, chatID, sessionK
return nil, false
}
func (mb *MessageBus) Stats() MessageBusStats {
if mb == nil {
return MessageBusStats{}
}
return MessageBusStats{
Inbound: snapshotStreamStats(mb.inbound, &mb.inboundStats),
Outbound: snapshotStreamStats(mb.outbound, &mb.outboundStats),
OutboundMedia: snapshotStreamStats(mb.outboundMedia, &mb.mediaStats),
AudioChunks: snapshotStreamStats(mb.audioChunks, &mb.audioStats),
VoiceControls: snapshotStreamStats(mb.voiceControls, &mb.voiceStats),
}
}
// HealthCheck returns a snapshot of queue depths and cumulative drop counts
// across all streams. It always reports ok=true: backpressure-induced drops are
// reflected in the message string for telemetry but do not affect the boolean.
// Callers that need readiness semantics (e.g. returning 503 when drops occur)
// should inspect Stats() directly and apply their own threshold logic.
func (mb *MessageBus) HealthCheck() (bool, string) {
stats := mb.Stats()
totalDropped := stats.Inbound.DroppedTotal +
stats.Outbound.DroppedTotal +
stats.OutboundMedia.DroppedTotal +
stats.AudioChunks.DroppedTotal +
stats.VoiceControls.DroppedTotal
message := fmt.Sprintf(
"in=%d/%d out=%d/%d media=%d/%d audio=%d/%d voice=%d/%d dropped=%d",
stats.Inbound.Depth,
stats.Inbound.Capacity,
stats.Outbound.Depth,
stats.Outbound.Capacity,
stats.OutboundMedia.Depth,
stats.OutboundMedia.Capacity,
stats.AudioChunks.Depth,
stats.AudioChunks.Capacity,
stats.VoiceControls.Depth,
stats.VoiceControls.Capacity,
totalDropped,
)
return true, message
}
func snapshotStreamStats[T any](ch chan T, stats *streamStats) StreamStats {
snapshot := StreamStats{
DroppedTotal: stats.dropped.Load(),
}
if ch != nil {
snapshot.Depth = len(ch)
snapshot.Capacity = cap(ch)
}
if unixNano := stats.lastDropped.Load(); unixNano > 0 {
snapshot.LastDroppedAt = time.Unix(0, unixNano)
}
if waitNanos := stats.lastWaitNanos.Load(); waitNanos > 0 {
wait := time.Duration(waitNanos)
snapshot.LastDropWait = wait.String()
snapshot.LastDropWaitMillis = wait.Milliseconds()
}
return snapshot
}
func (mb *MessageBus) Close() {
mb.closeOnce.Do(func() {
mb.publishCloseEvent(runtimeevents.KindBusCloseStarted, 0)
// notify all blocked publishers to exit
close(mb.done)
// because every publisher will check mb.closed before acquiring wg
// so we can be sure that new publishers will not be added new messages after this point
mb.publishMu.Lock()
mb.closed.Store(true)
close(mb.done)
mb.publishMu.Unlock()
// wait for all ongoing Publish calls to finish, ensuring all messages have been sent to channels or exited
mb.wg.Wait()

147
pkg/bus/bus_test.go
View file

@ -2,6 +2,7 @@ package bus
import (
"context"
"errors"
"sync"
"testing"
"time"
@ -459,6 +460,84 @@ func TestPublishAudioChunkSubscribe(t *testing.T) {
}
}
func TestPublishAudioChunk_BackpressureDropPublishesRuntimeEvent(t *testing.T) {
eventBus := runtimeevents.NewBus()
defer func() {
if err := eventBus.Close(); err != nil {
t.Errorf("event bus close failed: %v", err)
}
}()
_, eventsCh, err := eventBus.Channel().OfKind(runtimeevents.KindBusMessageDropped).SubscribeChan(
t.Context(),
runtimeevents.SubscribeOptions{Name: "bus-drop-events", Buffer: 1},
)
if err != nil {
t.Fatalf("SubscribeChan failed: %v", err)
}
mb := NewMessageBus()
defer mb.Close()
mb.SetEventPublisher(eventBus)
for i := range defaultBusBufferSize * 4 {
if pubErr := mb.PublishAudioChunk(context.Background(), AudioChunk{
SessionID: "voice-1",
SpeakerID: "speaker-1",
ChatID: "chat-1",
Channel: "discord",
Sequence: uint64(i),
Format: "opus",
Data: []byte{0x01},
}); pubErr != nil {
t.Fatalf("fill failed at %d: %v", i, pubErr)
}
}
err = mb.PublishAudioChunk(context.Background(), AudioChunk{
SessionID: "voice-1",
SpeakerID: "speaker-1",
ChatID: "chat-1",
Channel: "discord",
Sequence: 999,
Format: "opus",
Data: []byte{0x01},
})
if !errors.Is(err, ErrBusBackpressure) {
t.Fatalf("PublishAudioChunk() error = %v, want %v", err, ErrBusBackpressure)
}
evt := receiveBusRuntimeEvent(t, eventsCh)
if evt.Kind != runtimeevents.KindBusMessageDropped ||
evt.Source.Name != "audio_chunk" ||
evt.Severity != runtimeevents.SeverityWarn {
t.Fatalf("drop event = %+v", evt)
}
if evt.Scope.Channel != "discord" || evt.Scope.ChatID != "chat-1" {
t.Fatalf("drop event scope = %+v", evt.Scope)
}
if evt.Attrs["stream"] != "audio_chunk" ||
evt.Attrs["reason"] != "queue_full_timeout" ||
evt.Attrs["wait_ms"] != defaultAudioPublishTimeout.Milliseconds() ||
evt.Attrs["queue_depth"] != defaultBusBufferSize*4 ||
evt.Attrs["queue_capacity"] != defaultBusBufferSize*4 ||
evt.Attrs["dropped_total"] != uint64(1) {
t.Fatalf("drop event attrs = %#v", evt.Attrs)
}
stats := mb.Stats()
if stats.AudioChunks.DroppedTotal != 1 {
t.Fatalf("AudioChunks dropped = %d, want 1", stats.AudioChunks.DroppedTotal)
}
if stats.AudioChunks.Depth != defaultBusBufferSize*4 {
t.Fatalf("AudioChunks depth = %d, want %d", stats.AudioChunks.Depth, defaultBusBufferSize*4)
}
wantWaitMS := defaultAudioPublishTimeout.Milliseconds()
if stats.AudioChunks.LastDropWaitMillis != wantWaitMS {
t.Fatalf("AudioChunks last wait ms = %d, want %d", stats.AudioChunks.LastDropWaitMillis, wantWaitMS)
}
}
func TestPublishVoiceControlSubscribe(t *testing.T) {
mb := NewMessageBus()
defer mb.Close()
@ -728,6 +807,74 @@ func TestPublishInbound_FullBuffer(t *testing.T) {
}
}
// TestPublishInbound_FullBufferUsesBusBackpressureBudget exercises the generic
// publish() backpressure path directly (with a short 20ms timeout) rather than
// going through PublishInbound(). This avoids waiting for a long context timeout
// and keeps the test fast. Context validation and public-API wiring are covered
// by TestPublishInbound_FullBuffer and TestPublishInbound_ContextCancel.
func TestPublishInbound_FullBufferUsesBusBackpressureBudget(t *testing.T) {
mb := NewMessageBus()
defer mb.Close()
ch := make(chan InboundMessage, 1)
ch <- InboundMessage{Content: "fill"}
scope := runtimeevents.Scope{Channel: "test", ChatID: "chat-overflow"}
err := publish(context.Background(), mb, ch, InboundMessage{Content: "overflow"}, publishPolicy{
stream: "inbound",
timeout: 20 * time.Millisecond,
}, &mb.inboundStats, scope)
if !errors.Is(err, ErrBusBackpressure) {
t.Fatalf("publish() error = %v, want %v", err, ErrBusBackpressure)
}
stats := mb.Stats()
if stats.Inbound.DroppedTotal != 1 {
t.Fatalf("Inbound dropped = %d, want 1", stats.Inbound.DroppedTotal)
}
if stats.Inbound.LastDropWaitMillis != 20 {
t.Fatalf("Inbound last wait ms = %d, want 20", stats.Inbound.LastDropWaitMillis)
}
}
func TestMessageBusHealthCheckIncludesQueueDepthAndDrops(t *testing.T) {
mb := NewMessageBus()
defer mb.Close()
ok, msg := mb.HealthCheck()
if !ok {
t.Fatal("HealthCheck should remain ok for backpressure telemetry")
}
if msg == "" {
t.Fatal("HealthCheck message should not be empty")
}
for i := range cap(mb.audioChunks) {
if err := mb.PublishAudioChunk(context.Background(), AudioChunk{
Channel: "discord",
ChatID: "voice-room",
Sequence: uint64(i),
Data: []byte("fill"),
}); err != nil {
t.Fatalf("fill audio buffer at %d: %v", i, err)
}
}
_ = mb.PublishAudioChunk(context.Background(), AudioChunk{
Channel: "discord",
ChatID: "voice-room",
Sequence: 999,
Data: []byte("overflow"),
})
stats := mb.Stats()
if stats.AudioChunks.Depth != cap(mb.audioChunks) {
t.Fatalf("audio depth = %d, want %d", stats.AudioChunks.Depth, cap(mb.audioChunks))
}
if stats.AudioChunks.DroppedTotal != 1 {
t.Fatalf("audio dropped = %d, want 1", stats.AudioChunks.DroppedTotal)
}
}
func TestCloseIdempotent(t *testing.T) {
mb := NewMessageBus()

51
pkg/bus/events.go
View file

@ -1,6 +1,8 @@
package bus
import (
"time"
runtimeevents "github.com/sipeed/picoclaw/pkg/events"
)
@ -13,6 +15,15 @@ type busClosePayload struct {
Drained int `json:"drained,omitempty"`
}
type busMessageDroppedPayload struct {
Stream string `json:"stream"`
Reason string `json:"reason"`
WaitMS int64 `json:"wait_ms"`
QueueDepth int `json:"queue_depth"`
QueueCap int `json:"queue_capacity"`
DroppedTotal uint64 `json:"dropped_total"`
}
func (mb *MessageBus) publishFailure(stream string, scope runtimeevents.Scope, err error) {
if mb == nil || err == nil {
return
@ -38,6 +49,46 @@ func (mb *MessageBus) publishFailure(stream string, scope runtimeevents.Scope, e
})
}
func (mb *MessageBus) publishDrop(
stream string,
scope runtimeevents.Scope,
reason string,
wait time.Duration,
queueDepth, queueCap int,
droppedTotal uint64,
) {
if mb == nil {
return
}
publisher, ok := mb.eventPublisher.Load().(EventPublisher)
if !ok || publisher == nil {
return
}
publisher.PublishNonBlocking(runtimeevents.Event{
Kind: runtimeevents.KindBusMessageDropped,
Source: runtimeevents.Source{Component: "bus", Name: stream},
Scope: scope,
Severity: runtimeevents.SeverityWarn,
Payload: busMessageDroppedPayload{
Stream: stream,
Reason: reason,
WaitMS: wait.Milliseconds(),
QueueDepth: queueDepth,
QueueCap: queueCap,
DroppedTotal: droppedTotal,
},
Attrs: map[string]any{
"stream": stream,
"reason": reason,
"wait_ms": wait.Milliseconds(),
"queue_depth": queueDepth,
"queue_capacity": queueCap,
"dropped_total": droppedTotal,
},
})
}
func (mb *MessageBus) publishCloseEvent(kind runtimeevents.Kind, drained int) {
if mb == nil {
return

8
pkg/channels/discord/voice.go
View file

@ -298,9 +298,11 @@ func (c *DiscordChannel) receiveVoice(vc *discordgo.VoiceConnection, guildID str
Data: p.Opus,
}
ctx, cancel := context.WithTimeout(c.ctx, 100*time.Millisecond)
err := c.bus.PublishAudioChunk(ctx, chunk)
cancel()
// Pass the parent context directly; the bus applies its own
// audio drop budget internally.
// A 100ms caller-level timeout would fire before the bus's
// backpressure handling, masking drop counters and events.
err := c.bus.PublishAudioChunk(c.ctx, chunk)
if err != nil {
logger.ErrorCF("discord", "Failed to publish audio chunk", map[string]any{
"guild": guildID,

4
pkg/events/kind.go
View file

@ -68,6 +68,9 @@ const (
// KindBusPublishFailed is emitted when message bus publish fails.
KindBusPublishFailed Kind = "bus.publish.failed"
// KindBusMessageDropped is emitted when a message is dropped due to
// backpressure (channel full for longer than the drop budget).
KindBusMessageDropped Kind = "bus.message.dropped"
// KindBusCloseStarted is emitted when message bus close starts.
KindBusCloseStarted Kind = "bus.close.started"
// KindBusCloseCompleted is emitted when message bus close completes.
@ -133,6 +136,7 @@ var knownKinds = []Kind{
KindChannelMessageOutboundFailed,
KindChannelRateLimited,
KindBusPublishFailed,
KindBusMessageDropped,
KindBusCloseStarted,
KindBusCloseCompleted,
KindBusCloseDrained,