Skip to content

server: load initial source asynchronously so /api/ping is responsive immediately #2829

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
trungutt wants to merge 1 commit into from
+143 −12
Closed

server: load initial source asynchronously so /api/ping is responsive immediately #2829

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
30 changes: 28 additions & 2 deletions pkg/server/source_loader.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,19 @@ type sourceLoader struct {
inner config.Source
refreshInterval time.Duration

// ready is closed when the first load attempt has completed (whether it
// produced data, an error, or was cancelled). Read blocks on this so
// callers never observe the pre-load (nil, nil) state, preserving the
// contract of the previous synchronous-constructor implementation.
ready chan struct{}

mu sync.RWMutex
data []byte
err error
}

// NewSourceLoader creates a new source loader that caches and periodically refreshes a config source.
// The initial load runs in a goroutine; callers of Read block until the first load completes.
func NewSourceLoader(ctx context.Context, inner config.Source, refreshInterval time.Duration) *sourceLoader {
return newSourceLoader(ctx, inner, refreshInterval)
}
Expand All @@ -27,9 +34,19 @@ func newSourceLoader(ctx context.Context, inner config.Source, refreshInterval t
sl := &sourceLoader{
inner: inner,
refreshInterval: refreshInterval,
ready: make(chan struct{}),
}

sl.load(ctx)
// Run the initial load asynchronously so the constructor returns
// immediately. This lets the HTTP server start serving non-source
// endpoints (e.g. /api/ping, /health) while a potentially slow OCI
// pull or HTTPS GET is still in flight. Endpoints that need the
// source content (getAgents, createSession, ...) block in Read
// until this load finishes.
go func() {
defer close(sl.ready)
sl.load(ctx)
}()

if refreshInterval > 0 {
go sl.refreshLoop(ctx)
Expand All @@ -46,7 +63,16 @@ func (sl *sourceLoader) ParentDir() string {
return sl.inner.ParentDir()
}

func (sl *sourceLoader) Read(_ context.Context) ([]byte, error) {
// Read returns the latest cached data, blocking until the first load attempt
// has completed or ctx is cancelled. Once the first load has finished, Read
// returns the cached data/err snapshot without further blocking; the
// background refreshLoop is responsible for keeping the snapshot fresh.
func (sl *sourceLoader) Read(ctx context.Context) ([]byte, error) {
select {
Copy link
Copy Markdown

@docker-agent docker-agent May 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Non-deterministic Read result when context is already cancelled after initial load

When the ready channel is already closed (initial load complete) and the caller's ctx is already done, both select cases are simultaneously ready and Go picks one at random (~50/50). This means a caller with a cancelled request context can receive (nil, context.Canceled) even though valid cached data sits in sl.data.

The previous implementation declared Read(_ context.Context) — the context was fully discarded, so cancelled contexts never caused data-access failures. This PR changes that contract: a request context cancelled just after the initial load completes (e.g. a slow client dropping its HTTP connection) will now non-deterministically return an error instead of the cached data.

Minimal fix — check ready first with a non-blocking select, and only then honour context cancellation:

func (sl *sourceLoader) Read(ctx context.Context) ([]byte, error) {
    select {
    case <-sl.ready:
        // fast path: already loaded, fall through regardless of ctx state
    default:
        // initial load still in flight: wait for it or ctx cancellation
        select {
        case <-sl.ready:
        case <-ctx.Done():
            return nil, ctx.Err()
        }
    }
    sl.mu.RLock()
    defer sl.mu.RUnlock()
    return sl.data, sl.err
}

This preserves the blocking-until-loaded contract while never sacrificing already-cached data to a racing cancellation.

case <-sl.ready:
case <-ctx.Done():
return nil, ctx.Err()
}
sl.mu.RLock()
defer sl.mu.RUnlock()
return sl.data, sl.err
Expand Down
125 changes: 115 additions & 10 deletions pkg/server/source_loader_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,8 @@ func TestSourceLoader_Read_WithRefreshInterval_BeforeExpiry(t *testing.T) {
refreshInterval := 100 * time.Millisecond
sl := newSourceLoader(ctx, inner, refreshInterval)

// Read should return cached data immediately
// Read blocks until the asynchronous initial load completes, then
// returns the cached data.
data, err := sl.Read(ctx)
require.NoError(t, err)
assert.Equal(t, []byte("test data"), data)
Expand Down Expand Up @@ -115,14 +116,12 @@ func TestSourceLoader_Read_Error(t *testing.T) {
ctx := t.Context()
sl := newSourceLoader(ctx, inner, 0)

// Initial load failed
assert.Equal(t, 1, inner.getReadCount())

// Read should return the error from initial load
// Read blocks on the initial load and returns its error.
data, err := sl.Read(ctx)
require.Error(t, err)
assert.Equal(t, expectedErr, err)
assert.Nil(t, data)
assert.Equal(t, 1, inner.getReadCount())
}

func TestSourceLoader_Read_DataChanges(t *testing.T) {
Expand Down Expand Up @@ -169,17 +168,19 @@ func TestSourceLoader_Read_ZeroRefreshInterval(t *testing.T) {
ctx := t.Context()
sl := newSourceLoader(ctx, inner, 0)

initialReadCount := inner.getReadCount()
// First Read blocks on and then drives the single initial load.
data, err := sl.Read(ctx)
require.NoError(t, err)
assert.Equal(t, []byte("test data"), data)
require.Equal(t, 1, inner.getReadCount())

// Multiple reads with zero refresh interval
// Subsequent reads must not trigger any further Source.Read calls.
for range 10 {
data, err := sl.Read(ctx)
require.NoError(t, err)
assert.Equal(t, []byte("test data"), data)
}

// All reads should return cached data from startup
assert.Equal(t, initialReadCount, inner.getReadCount())
assert.Equal(t, 1, inner.getReadCount())
}

func TestSourceLoader_SuccessThenError(t *testing.T) {
Expand Down Expand Up @@ -211,3 +212,107 @@ func TestSourceLoader_SuccessThenError(t *testing.T) {
assert.Equal(t, []byte("initial data"), data)
})
}

// slowSource blocks Read until the test releases it. It models a slow OCI
// pull or HTTPS GET that would previously have blocked NewSessionManager.
type slowSource struct {
name string
release chan struct{}
data []byte
}

func (s *slowSource) Name() string { return s.name }
func (s *slowSource) ParentDir() string { return "" }

func (s *slowSource) Read(ctx context.Context) ([]byte, error) {
select {
case <-s.release:
return s.data, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}

// TestSourceLoader_ConstructorReturnsBeforeInitialLoad is the regression test
// for the docker-agent-as-AI-backend startup hang: NewSourceLoader (and
// therefore NewSessionManager and server.New) must return immediately, even
// if the underlying Source.Read is slow. Endpoints that don't need source
// data (notably /api/ping) can then start serving while the load is still
// in flight.
func TestSourceLoader_ConstructorReturnsBeforeInitialLoad(t *testing.T) {
t.Parallel()
inner := &slowSource{
name: "slow.yaml",
release: make(chan struct{}),
data: []byte("loaded data"),
}
ctx := t.Context()

constructorDone := make(chan struct{})
var sl *sourceLoader
go func() {
sl = newSourceLoader(ctx, inner, 0)
close(constructorDone)
}()

select {
case <-constructorDone:
case <-time.After(time.Second):
t.Fatal("newSourceLoader did not return while the source was still loading")
}

// Read should block until we release the underlying source.
readDone := make(chan struct {
data []byte
err error
})
go func() {
data, err := sl.Read(ctx)
readDone <- struct {
data []byte
err error
}{data, err}
}()

select {
case <-readDone:
t.Fatal("Read returned before the initial load completed")
case <-time.After(50 * time.Millisecond):
// Expected: Read is waiting on sl.ready
}

close(inner.release)

select {
case result := <-readDone:
require.NoError(t, result.err)
assert.Equal(t, []byte("loaded data"), result.data)
case <-time.After(time.Second):
t.Fatal("Read did not unblock after the initial load completed")
}
}

// TestSourceLoader_ReadCancelledByContext verifies that a caller of Read
// can give up if the initial load is taking too long, without leaking the
// loader's background goroutine.
func TestSourceLoader_ReadCancelledByContext(t *testing.T) {
t.Parallel()
inner := &slowSource{
name: "slow.yaml",
release: make(chan struct{}),
data: []byte("loaded data"),
}
loaderCtx, cancelLoader := context.WithCancel(t.Context())
defer cancelLoader()
sl := newSourceLoader(loaderCtx, inner, 0)

readCtx, cancelRead := context.WithCancel(t.Context())
cancelRead()

data, err := sl.Read(readCtx)
require.ErrorIs(t, err, context.Canceled)
assert.Nil(t, data)

// Let the loader's goroutine finish so we don't leak it across tests.
close(inner.release)
}
Loading