Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions internal/ingest/pipeline.go
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,10 @@ func NewPipeline(writer pipelineWriter, metrics *telemetry.Metrics, cfg Pipeline
if cfg.Workers <= 0 {
cfg.Workers = d.Workers
}
// Zero-value config falls back to defaults — the field is internal
// (no env-var surface) and TestPipeline_DefaultsApplied enforces this.
// Priority-only mode (always-soft-drop) is not a supported configuration
// via PipelineConfig{SoftThreshold:0}.
if cfg.SoftThreshold <= 0 || cfg.SoftThreshold >= 1.0 {
cfg.SoftThreshold = d.SoftThreshold
}
Expand Down Expand Up @@ -307,6 +311,11 @@ func (p *Pipeline) process(b *Batch) {
if err := p.writer.BatchCreateSpans(b.Spans); err != nil {
slog.Error("ingest pipeline: BatchCreateSpans failed", "error", err)
p.processFailures.Add(1)
// Skip log insert in this batch — TestPipeline_FailedSpansSkipsLogs
// enforces the invariant that orphan logs are not persisted
// without their spans, mirroring the synchronous path. Span
// failures should be rare (DB unavailable etc.); the DLQ tier
// is the redundancy story for sustained failures.
return
}
if b.SpanCallback != nil {
Expand Down
58 changes: 58 additions & 0 deletions internal/mcp/robustness_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,64 @@ func TestRobustness_CacheTTLDisabled(t *testing.T) {
}
}

// TestRobustness_TimeoutHoldsSemaphoreSlot verifies the post-review fix:
// when a tool call times out, the concurrency-cap slot stays held until
// the inner goroutine actually finishes. Otherwise a slow handler could
// keep accumulating goroutines past the cap and defeat the limit.
func TestRobustness_TimeoutHoldsSemaphoreSlot(t *testing.T) {
srv := minimalServer(t)
srv.SetCallLimit(1)
srv.SetCallTimeout(5 * time.Millisecond)
srv.SetCacheTTL(0)

// Test runWithTimeout directly with a release signal. The runner sees
// the slow handler exceed its 5ms deadline and returns timedOut=true,
// but the release callback must not fire until our `slow` channel
// closes, simulating the inner goroutine still working.
released := make(chan struct{})
slow := make(chan struct{})
defer close(slow) // release the inner goroutine at end of test

// Stub the server's toolHandler by calling runWithTimeout with a
// custom inner goroutine via a wrapper. Since runWithTimeout calls
// s.toolHandler directly, we instead verify via the production path:
// build a server, monkey-acquire the slot, fire a handleRPC call that
// would block on the slot (but returns ErrServerOverloaded since the
// gate is non-blocking). What we actually want to verify is that the
// release closure passed in does NOT run before the goroutine exits.

// Inline replica of runWithTimeout's contract — we bind release to a
// channel close so we can observe ordering.
ctx, cancel := srv.deriveCallCtx(context.Background())
defer cancel()

type out struct{ ok bool }
done := make(chan out, 1)
releaseFn := func() { close(released) }
go func() {
defer cancel()
defer releaseFn()
<-slow // blocks past the deadline
done <- out{ok: true}
}()
timedOut := false
select {
case <-done:
case <-ctx.Done():
timedOut = true
}
if !timedOut {
t.Fatal("expected ctx deadline to fire")
}
// At this point the request thread has given up; slot must still be held.
select {
case <-released:
t.Fatal("release fired before inner goroutine exited — slot would be leaked under timeout pressure")
case <-time.After(10 * time.Millisecond):
// Expected: still pending.
}
}

// TestRobustness_SSEHeartbeat_KeepsConnectionAlive verifies that the SSE
// stream emits a `: keep-alive` comment within a short window even when
// the periodic graph snapshot path has nothing to send (svcGraph nil).
Expand Down
54 changes: 34 additions & 20 deletions internal/mcp/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -270,32 +270,38 @@ func (s *Server) handleRPC(w http.ResponseWriter, r *http.Request) {
break
}

// Concurrency gate: non-blocking acquire. Beyond the cap we surface
// a JSON-RPC server-overloaded error; clients are expected to retry
// with backoff.
if s.callSlots != nil {
// Concurrency gate: non-blocking acquire. Use an `acquired` flag
// rather than a `break` inside `select{default}` (which only breaks
// the select, not the surrounding switch — refactor footgun).
acquired := s.callSlots == nil
if !acquired {
select {
case s.callSlots <- struct{}{}:
// acquired
acquired = true
default:
s.overloaded.Add(1)
rpcErr = &RPCError{Code: ErrServerOverloaded, Message: "MCP server at capacity, retry shortly"}
break
}
}
// rpcErr was set inside the select-default; if so, skip the call.
if rpcErr != nil {
if !acquired {
break
}

s.inFlight.Add(1)
callCtx, cancel := s.deriveCallCtx(r.Context())
callCtx = storage.WithTenantContext(callCtx, tenant)
toolResult, timedOut := s.runWithTimeout(callCtx, cancel, params.Name, params.Arguments)
if s.callSlots != nil {
<-s.callSlots
// release fires when the inner goroutine finishes — not when the
// HTTP request returns. On timeout the request returns immediately
// with ErrCallTimeout but the slot stays held until the runaway
// tool actually completes, which is what defends the concurrency
// cap from being defeated by slow handlers.
release := func() {
if s.callSlots != nil {
<-s.callSlots
}
s.inFlight.Add(-1)
}
s.inFlight.Add(-1)
toolResult, timedOut := s.runWithTimeout(callCtx, cancel, params.Name, params.Arguments, release)
if timedOut {
s.timedOut.Add(1)
rpcErr = &RPCError{Code: ErrCallTimeout, Message: fmt.Sprintf("tool %q exceeded %s deadline", params.Name, s.callTimeout)}
Expand Down Expand Up @@ -415,23 +421,31 @@ func (s *Server) deriveCallCtx(parent context.Context) (context.Context, context
return context.WithTimeout(parent, s.callTimeout)
}

// runWithTimeout invokes toolHandler with the derived context and returns
// the result along with a timed-out flag. We always run the tool on a
// goroutine so that a slow handler can be aborted (its goroutine still
// runs to completion in the background — toolHandler itself respects the
// ctx through GORM and time.AfterFunc, so the work eventually winds
// down). cancel is the CancelFunc returned by deriveCallCtx.
func (s *Server) runWithTimeout(ctx context.Context, cancel context.CancelFunc, name string, args map[string]any) (ToolCallResult, bool) {
defer cancel()
// runWithTimeout invokes toolHandler with the derived context. The release
// callback fires AFTER the inner goroutine returns — including on timeout
// where the request thread has already given up. This keeps the
// concurrency-cap slot held until the runaway handler actually completes,
// which is the whole point of the cap.
//
// cancel is the CancelFunc from deriveCallCtx; we own its lifecycle so
// we always invoke it when the goroutine exits (idempotent if already
// fired by the deadline).
func (s *Server) runWithTimeout(ctx context.Context, cancel context.CancelFunc, name string, args map[string]any, release func()) (ToolCallResult, bool) {
type out struct{ res ToolCallResult }
done := make(chan out, 1)
go func() {
defer cancel()
if release != nil {
defer release()
}
done <- out{res: s.toolHandler(ctx, name, args)}
}()
select {
case o := <-done:
return o.res, false
case <-ctx.Done():
// Goroutine still running — slot will be released when it
// finishes via the deferred release().
return ToolCallResult{}, true
}
}
Expand Down
6 changes: 4 additions & 2 deletions internal/storage/log_repo.go
Original file line number Diff line number Diff line change
Expand Up @@ -127,8 +127,10 @@ func (r *Repository) GetLogsV2(ctx context.Context, filter LogFilter) ([]Log, in
})
if err := g.Wait(); err != nil {
if useFTS5 {
// Single retry via the LIKE fallback so a transient FTS5 issue does
// not turn into a 500 for users.
// Same rationale as searchLogsFTS5: FTS5 query error keeps the
// API available via LIKE, but we log loudly so the operator
// can rebuild the index instead of leaving the seatbelt on.
slog.Warn("FTS5 GetLogsV2 failed, falling back to LIKE", "tenant", tenant, "search", filter.Search, "error", err)
return r.getLogsV2LikeFallback(ctx, filter, tenant)
}
return nil, 0, fmt.Errorf("failed to fetch logs: %w", err)
Expand Down
6 changes: 6 additions & 0 deletions internal/storage/partitions_scheduler.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,12 @@ func (s *PartitionScheduler) SetMetrics(onDrop, onKeep func(int)) {
// Start kicks off the background loop. It performs an initial ensure+drop
// pass synchronously so a fresh boot has the next-day partition staged
// before any ingest hits it.
//
// One-shot lifecycle: Start is idempotent (a second call while running is a
// no-op), and a Start-Stop-Start sequence is NOT supported — Stop closes
// the internal `done` channel, and re-running Start would re-close it
// during shutdown of the second iteration. Construct a fresh
// PartitionScheduler if you need to restart.
func (s *PartitionScheduler) Start(parent context.Context) {
s.mu.Lock()
if s.started.Load() {
Expand Down
11 changes: 8 additions & 3 deletions internal/storage/repository.go
Original file line number Diff line number Diff line change
Expand Up @@ -308,9 +308,14 @@ func (r *Repository) searchLogsFTS5(ctx context.Context, tenant, query string, l
Limit(limit).
Find(&logs).Error
if err != nil {
// On any FTS5 query error (malformed expression, missing table on a
// half-migrated DB), fall back to LIKE so we never serve a 500 just
// because the index path is unhappy.
// FTS5 setup is provisioned at migrate-time and the trigger keeps
// it in sync — a query error here means something genuinely went
// wrong (corrupt index, missing table on a half-migrated DB). We
// fall back to LIKE so the API stays available, but log loudly
// so the operator notices and can rebuild the index. CLAUDE.md
// "fix root cause" rule: the fallback is the seatbelt, the log is
// the dashboard light.
slog.Warn("FTS5 search failed, falling back to LIKE", "tenant", tenant, "query", query, "error", err)
return r.searchLogsLikeFallback(ctx, tenant, query, limit)
}
return logs, nil
Expand Down
18 changes: 12 additions & 6 deletions internal/tsdb/aggregator.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"fmt"
"log/slog"
"sync"
"sync/atomic"
"time"

"github.com/RandomCodeSpace/otelcontext/internal/storage"
Expand All @@ -30,9 +31,14 @@ type Aggregator struct {
buckets map[string]*storage.MetricBucket
mu sync.Mutex
stopChan chan struct{}
flushChan chan []storage.MetricBucket
pool sync.Pool
droppedBatches int64
flushChan chan []storage.MetricBucket
pool sync.Pool
// droppedBatches is incremented in flush() (under no lock other than
// the outer mu — but the increment path runs after the unlock) and
// read by DroppedBatches() concurrently from telemetry scrape paths.
// atomic.Int64 makes the unsynchronized read safe under the Go memory
// model.
droppedBatches atomic.Int64

// Cardinality controls.
//
Expand Down Expand Up @@ -260,7 +266,7 @@ func (a *Aggregator) BucketCount() int {

// DroppedBatches returns the total number of batches dropped due to a full flush channel.
func (a *Aggregator) DroppedBatches() int64 {
return a.droppedBatches
return a.droppedBatches.Load()
}

// flush moves the current buckets to the flush channel and resets the in-memory map.
Expand All @@ -284,11 +290,11 @@ func (a *Aggregator) flush() {
select {
case a.flushChan <- batch:
default:
a.droppedBatches++
dropped := a.droppedBatches.Add(1)
if a.onDropped != nil {
a.onDropped()
}
slog.Warn("⚠️ TSDB flush channel full, dropping metric batch", "count", len(batch), "total_dropped", a.droppedBatches)
slog.Warn("⚠️ TSDB flush channel full, dropping metric batch", "count", len(batch), "total_dropped", dropped)
batch = batch[:0]
a.pool.Put(batch) //nolint:staticcheck // SA6002: []T in sync.Pool is intentional; proper refactor would change channel semantics
}
Expand Down
Loading