wdclient: keep a shared volume lookup alive when its first caller leaves (#12579)

* wdclient: keep a shared volume lookup alive when its first caller leaves

Concurrent misses for the same volumes share one provider lookup, but it
ran on the first caller's context, so that caller's cancellation failed
every request waiting on it. The lookup now runs on a context no single
caller can cancel, and each caller stops waiting when its own context ends.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* wdclient: wait for the second caller to join before canceling the first

A fixed sleep did not prove the second caller had joined the shared lookup:
a late caller could hit the filled cache and pass without waiting. The test
now waits until the caller reads its own Done channel, which happens only
after DoChan registered it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

weed/wdclient/vidmap_client.go

@@ -209,7 +209,7 @@ func (vc *vidMapClient) LookupVolumeIdsWithFallback(ctx context.Context, volumeI
 	sort.Strings(needsLookup)
 	batchKey := strings.Join(needsLookup, ",")
 
-	sfResult, err, _ := vc.vidLookupGroup.Do(batchKey, func() (interface{}, error) {
+	sfResult, err := vc.doSharedLookup(ctx, batchKey, func(ctx context.Context) (interface{}, error) {
 		// Double-check cache for volumes that might have been populated while waiting
 		stillNeedLookup := make([]string, 0, len(needsLookup))
 		batchResult := make(map[string][]Location)
@@ -280,6 +280,21 @@ func (vc *vidMapClient) LookupVolumeIdsWithFallback(ctx context.Context, volumeI
 	return result, errors.Join(lookupErrors...)
 }
 
+// doSharedLookup runs fn once for all concurrent callers of key. fn gets a
+// context no single caller can cancel, since one caller leaving must not fail
+// the others; each caller stops waiting when its own context ends.
+func (vc *vidMapClient) doSharedLookup(ctx context.Context, key string, fn func(context.Context) (interface{}, error)) (interface{}, error) {
+	resultCh := vc.vidLookupGroup.DoChan(key, func() (interface{}, error) {
+		return fn(context.WithoutCancel(ctx))
+	})
+	select {
+	case res := <-resultCh:
+		return res.Val, res.Err
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
 // LookupVolumeIdsAuthoritative bypasses the asynchronously maintained vid map.
 // Use it when a stale positive result is unsafe, such as before serving cache
 // payload from a Volume that the master may have retired. The provider still

weed/wdclient/vidmap_shared_lookup_test.go

@@ -0,0 +1,95 @@
+package wdclient
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+// blockingLookupProvider answers once release is closed, or fails with the
+// context error if its context ends first, as the gRPC providers do.
+type blockingLookupProvider struct {
+	calls   atomic.Int32
+	started chan struct{}
+	once    sync.Once
+	release chan struct{}
+}
+
+func (p *blockingLookupProvider) LookupVolumeIds(ctx context.Context, volumeIds []string) (map[string][]Location, error) {
+	p.calls.Add(1)
+	p.once.Do(func() { close(p.started) })
+	select {
+	case <-p.release:
+		return map[string][]Location{"7": {{Url: "10.0.0.7:8080"}}}, nil
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
+// joinSignalCtx closes joined the first time Done is read. doSharedLookup reads
+// it only after DoChan has registered the caller with the flight.
+type joinSignalCtx struct {
+	context.Context
+	once   sync.Once
+	joined chan struct{}
+}
+
+func (c *joinSignalCtx) Done() <-chan struct{} {
+	c.once.Do(func() { close(c.joined) })
+	return c.Context.Done()
+}
+
+func TestSharedLookupSurvivesFirstCallerCancel(t *testing.T) {
+	provider := &blockingLookupProvider{started: make(chan struct{}), release: make(chan struct{})}
+	vc := newVidMapClient(provider, "", DefaultVidMapCacheSize)
+
+	firstCtx, cancelFirst := context.WithCancel(context.Background())
+	firstErr := make(chan error, 1)
+	go func() {
+		_, err := vc.LookupVolumeIdsWithFallback(firstCtx, []string{"7"})
+		firstErr <- err
+	}()
+	<-provider.started
+
+	type lookupResult struct {
+		locations map[string][]Location
+		err       error
+	}
+	second := make(chan lookupResult, 1)
+	secondCtx := &joinSignalCtx{Context: context.Background(), joined: make(chan struct{})}
+	go func() {
+		locations, err := vc.LookupVolumeIdsWithFallback(secondCtx, []string{"7"})
+		second <- lookupResult{locations, err}
+	}()
+	select {
+	case <-secondCtx.joined:
+	case <-time.After(5 * time.Second):
+		t.Fatal("second caller never joined the shared lookup on its own context")
+	}
+
+	cancelFirst()
+	select {
+	case err := <-firstErr:
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("canceled caller got %v, want context.Canceled", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("canceled caller kept waiting on the shared lookup")
+	}
+
+	close(provider.release)
+	select {
+	case res := <-second:
+		if res.err != nil || len(res.locations["7"]) != 1 {
+			t.Fatalf("waiter got %v, err %v after another caller canceled", res.locations["7"], res.err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("waiter never got the shared lookup result")
+	}
+	if calls := provider.calls.Load(); calls != 1 {
+		t.Fatalf("provider called %d times, want one shared lookup", calls)
+	}
+}