Skip to content

Commit 4a51143

Browse files
committed
feat(opensearchtransport): make DiscoverNodes blocking, add *bool DiscoverNodesOnStart
DiscoverNodes now waits for an in-flight discovery to complete (or for the context to be cancelled) instead of returning nil immediately. This lets callers block until topology data is available after client construction. - Auto-enable DiscoverNodesOnStart when OPENSEARCH_GO_ROUTER=true and the caller did not set the field Signed-off-by: Sean Chittenden <sean.chittenden@crowdstrike.com>
1 parent aaa96fc commit 4a51143

5 files changed

Lines changed: 284 additions & 20 deletions

File tree

opensearch_internal_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ import (
5050

5151
var called int
5252

53+
func boolPtr(v bool) *bool { return &v }
54+
5355
var defaultRoundTripFunc = func(req *http.Request) (*http.Response, error) {
5456
response := &http.Response{Header: http.Header{}}
5557

@@ -78,8 +80,6 @@ type testReq struct {
7880
Headers http.Header
7981
}
8082

81-
func boolPtr(v bool) *bool { return &v }
82-
8383
func (r testReq) GetRequest(method string) (*http.Request, error) {
8484
if r.Error {
8585
return nil, fmt.Errorf("test error")

opensearchtransport/discovery.go

Lines changed: 73 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -252,34 +252,90 @@ func (m *_NodesMeta) formatFailures() string {
252252
return string(b)
253253
}
254254

255-
// DiscoverNodes reloads the client connections by fetching information from the cluster.
255+
// DiscoverNodes reloads the client connections by fetching information from
256+
// the cluster. If another discovery is already in progress, DiscoverNodes
257+
// blocks until that discovery completes (or ctx is cancelled) and returns
258+
// its result.
256259
func (c *Client) DiscoverNodes(ctx context.Context) error {
257-
// Bail out early if the context is already cancelled (e.g. client shutting down).
258260
if ctx.Err() != nil {
259261
return ctx.Err()
260262
}
261263

262-
// Prevent concurrent discovery operations
263-
c.mu.Lock()
264-
if c.mu.discoveryInProgress {
265-
c.mu.Unlock()
264+
c.discoverMu.Lock()
265+
266+
if c.discoverMu.inProgress {
267+
// Another goroutine is running discovery. Wait for it using
268+
// sync.Cond + context.AfterFunc so that context cancellation
269+
// wakes us even though Cond.Wait is not context-aware.
270+
stopf := context.AfterFunc(ctx, func() {
271+
c.discoverMu.Lock()
272+
defer c.discoverMu.Unlock()
273+
c.discoverMu.cond.Broadcast()
274+
})
275+
defer stopf()
276+
277+
for c.discoverMu.inProgress {
278+
c.discoverMu.cond.Wait()
279+
if ctx.Err() != nil {
280+
c.discoverMu.Unlock()
281+
return ctx.Err()
282+
}
283+
}
284+
err := c.discoverMu.lastErr
285+
c.discoverMu.Unlock()
286+
return err
287+
}
288+
289+
// We won the race: start discovery.
290+
// Lock is held — doDiscoverNodes takes ownership and releases it.
291+
return c.doDiscoverNodes(ctx)
292+
}
293+
294+
// tryDiscoverNodes attempts to start a discovery cycle. If discovery is
295+
// already in progress it returns nil immediately without waiting.
296+
//
297+
// This is used by the internal discoveryLoop, which must never block on
298+
// another discovery. It could be exported in the future if callers need
299+
// fire-and-forget semantics on the public API.
300+
func (c *Client) tryDiscoverNodes(ctx context.Context) error {
301+
if ctx.Err() != nil {
302+
return ctx.Err()
303+
}
304+
305+
c.discoverMu.Lock()
306+
if c.discoverMu.inProgress {
307+
c.discoverMu.Unlock()
266308
return nil
267309
}
268-
c.mu.discoveryInProgress = true
269-
c.mu.Unlock()
310+
// Lock is held — doDiscoverNodes takes ownership and releases it.
311+
return c.doDiscoverNodes(ctx)
312+
}
270313

314+
// doDiscoverNodes performs the discovery work.
315+
//
316+
// Called with c.discoverMu held. It sets inProgress = true, releases the
317+
// lock for I/O, then re-acquires it on completion to clear inProgress,
318+
// store the result, and wake any waiters.
319+
func (c *Client) doDiscoverNodes(ctx context.Context) error {
320+
c.discoverMu.inProgress = true
321+
c.discoverMu.Unlock()
322+
323+
var discoverErr error
271324
defer func() {
272-
c.mu.Lock()
273-
c.mu.discoveryInProgress = false
274-
c.mu.Unlock()
325+
c.discoverMu.Lock()
326+
c.discoverMu.inProgress = false
327+
c.discoverMu.lastErr = discoverErr
328+
c.discoverMu.cond.Broadcast()
329+
c.discoverMu.Unlock()
275330
}()
276331

277332
discovered, err := c.getNodesInfo(ctx)
278333
if err != nil {
279334
if dl := loadDebugLogger(); dl != nil {
280335
dl.Logf("Error getting nodes info: %s\n", err)
281336
}
282-
return fmt.Errorf("discovery: get nodes: %w", err)
337+
discoverErr = fmt.Errorf("discovery: get nodes: %w", err)
338+
return discoverErr
283339
}
284340

285341
c.mu.RLock()
@@ -289,11 +345,13 @@ func (c *Client) DiscoverNodes(ctx context.Context) error {
289345

290346
if isColdStart {
291347
if err := c.nodeDiscoveryAsyncStart(ctx, discovered); err != nil {
292-
return err
348+
discoverErr = err
349+
return discoverErr
293350
}
294351
} else {
295352
if err := c.nodeDiscovery(ctx, discovered); err != nil {
296-
return err
353+
discoverErr = err
354+
return discoverErr
297355
}
298356
}
299357

@@ -1217,7 +1275,7 @@ func (c *Client) discoveryLoop() {
12171275
switch {
12181276
case !now.Before(nextNodes):
12191277
// Full node + shard discovery.
1220-
c.DiscoverNodes(c.ctx) //nolint:errcheck // errors logged inside
1278+
c.tryDiscoverNodes(c.ctx) //nolint:errcheck // errors logged inside
12211279

12221280
nextNodes = time.Now().Add(c.discoverNodesInterval)
12231281
nextCat = time.Time{}

opensearchtransport/discovery_internal_test.go

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import (
4242
"slices"
4343
"strconv"
4444
"strings"
45+
"sync"
4546
"testing"
4647
"time"
4748

@@ -2259,3 +2260,198 @@ func TestGetNodesInfoNodesMeta(t *testing.T) {
22592260
})
22602261
}
22612262
}
2263+
2264+
// gatedNodesHandler returns a /_nodes/http handler that signals entered when
2265+
// the request arrives, then blocks until gate is closed before responding with
2266+
// a single data node. The handler also selects on t.Context().Done() so that
2267+
// test cleanup can unblock it.
2268+
func gatedNodesHandler(t *testing.T, entered chan<- struct{}, gate <-chan struct{}) http.HandlerFunc {
2269+
t.Helper()
2270+
return func(w http.ResponseWriter, _ *http.Request) {
2271+
select {
2272+
case entered <- struct{}{}:
2273+
default:
2274+
}
2275+
select {
2276+
case <-gate:
2277+
case <-t.Context().Done():
2278+
http.Error(w, "test context cancelled", http.StatusServiceUnavailable)
2279+
return
2280+
}
2281+
w.Header().Set("Content-Type", "application/json")
2282+
fmt.Fprint(w, `{
2283+
"_nodes":{"total":1,"successful":1,"failed":0},
2284+
"cluster_name":"test",
2285+
"nodes":{
2286+
"n1":{
2287+
"name":"n1",
2288+
"roles":["data","ingest"],
2289+
"http":{"publish_address":"127.0.0.1:9200"}
2290+
}
2291+
}
2292+
}`)
2293+
}
2294+
}
2295+
2296+
// gatedErrorNodesHandler returns a /_nodes/http handler that signals entered
2297+
// when the request arrives, then blocks until gate is closed before responding
2298+
// with an HTTP 503 to trigger a discovery error. The handler also selects on
2299+
// t.Context().Done() so that test cleanup can unblock it.
2300+
func gatedErrorNodesHandler(t *testing.T, entered chan<- struct{}, gate <-chan struct{}) http.HandlerFunc {
2301+
t.Helper()
2302+
return func(w http.ResponseWriter, _ *http.Request) {
2303+
select {
2304+
case entered <- struct{}{}:
2305+
default:
2306+
}
2307+
select {
2308+
case <-gate:
2309+
case <-t.Context().Done():
2310+
}
2311+
http.Error(w, "unavailable", http.StatusServiceUnavailable)
2312+
}
2313+
}
2314+
2315+
// newGatedDiscoverClient creates a transport Client wired to the given
2316+
// handler routes, with discoverMu.cond properly initialised.
2317+
func newGatedDiscoverClient(t *testing.T, routes mockhttp.HandlerMap) *Client {
2318+
t.Helper()
2319+
transport := mockhttp.NewTransportFromRoutes(t, routes)
2320+
u, _ := url.Parse("http://127.0.0.1:9200")
2321+
tp, err := New(Config{URLs: []*url.URL{u}, Transport: transport})
2322+
require.NoError(t, err)
2323+
tp.discoverMu.cond = sync.NewCond(&tp.discoverMu)
2324+
return tp
2325+
}
2326+
2327+
func TestDiscoverNodesBlocking(t *testing.T) {
2328+
entered := make(chan struct{}, 1)
2329+
gate := make(chan struct{})
2330+
2331+
routes := mockhttp.GetDefaultHandlers(t)
2332+
routes["/_nodes/http"] = gatedNodesHandler(t, entered, gate)
2333+
tp := newGatedDiscoverClient(t, routes)
2334+
2335+
// Goroutine A: start discovery (blocks in handler on gate).
2336+
var wg sync.WaitGroup
2337+
wg.Add(1)
2338+
go func() {
2339+
defer wg.Done()
2340+
tp.DiscoverNodes(t.Context())
2341+
}()
2342+
2343+
// Wait for handler to be entered — discovery is now in-flight.
2344+
<-entered
2345+
2346+
// Goroutine B: should block in DiscoverNodes until A finishes.
2347+
bDone := make(chan error, 1)
2348+
go func() {
2349+
bDone <- tp.DiscoverNodes(t.Context())
2350+
}()
2351+
2352+
// B should not have returned yet (gate still closed, A still blocked).
2353+
select {
2354+
case <-bDone:
2355+
t.Fatal("goroutine B returned before discovery finished")
2356+
default:
2357+
}
2358+
2359+
// Release A — A finishes, B wakes up.
2360+
close(gate)
2361+
wg.Wait()
2362+
2363+
err := <-bDone
2364+
require.NoError(t, err, "goroutine B should succeed after waiting")
2365+
}
2366+
2367+
func TestDiscoverNodesBlockingPropagatesError(t *testing.T) {
2368+
entered := make(chan struct{}, 1)
2369+
gate := make(chan struct{})
2370+
2371+
routes := mockhttp.GetDefaultHandlers(t)
2372+
routes["/_nodes/http"] = gatedErrorNodesHandler(t, entered, gate)
2373+
tp := newGatedDiscoverClient(t, routes)
2374+
2375+
// Goroutine A: start discovery that will fail.
2376+
aDone := make(chan error, 1)
2377+
go func() {
2378+
aDone <- tp.DiscoverNodes(t.Context())
2379+
}()
2380+
2381+
<-entered // A is in handler, discovery in-flight.
2382+
2383+
// Goroutine B: waits for A, receives the same error via lastErr.
2384+
bDone := make(chan error, 1)
2385+
go func() {
2386+
bDone <- tp.DiscoverNodes(t.Context())
2387+
}()
2388+
2389+
close(gate)
2390+
2391+
errA := <-aDone
2392+
errB := <-bDone
2393+
2394+
require.Error(t, errA, "runner should report discovery error")
2395+
require.Error(t, errB, "waiter should receive the same error")
2396+
require.Equal(t, errA, errB)
2397+
}
2398+
2399+
func TestDiscoverNodesBlockingContextCancel(t *testing.T) {
2400+
entered := make(chan struct{}, 1)
2401+
gate := make(chan struct{})
2402+
defer close(gate) // prevent goroutine leak
2403+
2404+
routes := mockhttp.GetDefaultHandlers(t)
2405+
routes["/_nodes/http"] = gatedNodesHandler(t, entered, gate)
2406+
tp := newGatedDiscoverClient(t, routes)
2407+
2408+
// Goroutine A: start slow discovery.
2409+
go func() {
2410+
tp.DiscoverNodes(t.Context())
2411+
}()
2412+
2413+
<-entered // A is in handler, discovery in-flight.
2414+
2415+
// B: call DiscoverNodes with an already-cancelled context.
2416+
ctx, cancel := context.WithCancel(t.Context())
2417+
cancel()
2418+
2419+
err := tp.DiscoverNodes(ctx)
2420+
require.ErrorIs(t, err, context.Canceled)
2421+
}
2422+
2423+
func TestTryDiscoverNodesNonBlocking(t *testing.T) {
2424+
entered := make(chan struct{}, 1)
2425+
gate := make(chan struct{})
2426+
defer close(gate) // prevent goroutine leak
2427+
2428+
routes := mockhttp.GetDefaultHandlers(t)
2429+
routes["/_nodes/http"] = gatedNodesHandler(t, entered, gate)
2430+
tp := newGatedDiscoverClient(t, routes)
2431+
2432+
// Start discovery in a goroutine.
2433+
go func() {
2434+
tp.DiscoverNodes(t.Context())
2435+
}()
2436+
2437+
<-entered // discovery is in-flight.
2438+
2439+
// tryDiscoverNodes should return nil immediately without blocking.
2440+
err := tp.tryDiscoverNodes(t.Context())
2441+
require.NoError(t, err)
2442+
}
2443+
2444+
func TestDiscoverNodesSequential(t *testing.T) {
2445+
routes := mockhttp.GetDefaultHandlersWithNodes(t, map[string][]string{
2446+
"node1": {"data", "ingest"},
2447+
})
2448+
tp := newGatedDiscoverClient(t, routes)
2449+
2450+
// First call succeeds.
2451+
err := tp.DiscoverNodes(t.Context())
2452+
require.NoError(t, err)
2453+
2454+
// Second call also succeeds (starts a fresh discovery).
2455+
err = tp.DiscoverNodes(t.Context())
2456+
require.NoError(t, err)
2457+
}

opensearchtransport/opensearchtransport.go

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -472,8 +472,17 @@ type Client struct {
472472

473473
mu struct {
474474
sync.RWMutex
475-
connectionPool ConnectionPool // Used for both single-node and multi-node
476-
discoveryInProgress bool // Prevents concurrent discovery operations
475+
connectionPool ConnectionPool // Used for both single-node and multi-node
476+
}
477+
478+
// discoverMu serializes discovery and lets callers block until an
479+
// in-flight discovery completes. Separate from mu so that
480+
// discovery waiters never contend with normal request-path readers.
481+
discoverMu struct {
482+
sync.Mutex
483+
cond *sync.Cond // signaled when inProgress transitions to false
484+
lastErr error // result of the most recent completed discovery
485+
inProgress bool // true while a discovery cycle is running
477486
}
478487
}
479488

@@ -845,6 +854,7 @@ func New(cfg Config) (*Client, error) {
845854
}
846855

847856
client.userAgent = initUserAgent()
857+
client.discoverMu.cond = sync.NewCond(&client.discoverMu)
848858

849859
// Parse discovery feature config from environment variable.
850860
// Controls which server calls are made during the discovery cycle.

opensearchtransport/pool_congestion_internal_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -891,7 +891,7 @@ type testDebugLogger struct{}
891891
func (l *testDebugLogger) Log(_ ...any) error { return nil }
892892
func (l *testDebugLogger) Logf(_ string, _ ...any) error { return nil }
893893

894-
// enableTestDebugLogger sets debugLogger to a no-op testDebugLogger exactly
894+
// enableTestDebugLogger sets the debug logger to a no-op testDebugLogger exactly
895895
// once for the lifetime of the test process. This avoids data races that
896896
// arise when individual tests save/restore the package-level global while
897897
// background goroutines from parallel tests are still reading it.

0 commit comments

Comments
 (0)