Merge pull request #365 from rest-sh/fix/release-qa-discovery-race

fix: wait for canceled discovery probes

Author: danielgtaylor

internal/cli/api_manage_test.go

@@ -699,6 +699,7 @@ func TestAPISyncDiscoverySendsAuthForShorthandBaseURL(t *testing.T) {
 }
 
 func TestAPISyncDiscoveryFollowsLinkWithShorthandBaseURL(t *testing.T) {
+	var mu sync.Mutex
 	var linkedSpecAuth string
 
 	c, _, _ := newAPISyncCLI(t, protectedAPI("api.example.com", "", bearerAuth("test-token")))
@@ -709,7 +710,9 @@ func TestAPISyncDiscoveryFollowsLinkWithShorthandBaseURL(t *testing.T) {
 			resp.Header.Set("Link", `</openapi.yaml>; rel="service-desc"`)
 			return resp, nil
 		case r.URL.Host == "api.example.com" && r.URL.Path == "/openapi.yaml":
+			mu.Lock()
 			linkedSpecAuth = r.Header.Get("Authorization")
+			mu.Unlock()
 			return jsonResponse(200, minimalOpenAPI), nil
 		default:
 			return textResponse(404, "text/plain", "not found", r), nil
@@ -719,6 +722,8 @@ func TestAPISyncDiscoveryFollowsLinkWithShorthandBaseURL(t *testing.T) {
 	if err := runProtectedAPISync(t, c); err != nil {
 		t.Fatalf("api sync: %v", err)
 	}
+	mu.Lock()
+	defer mu.Unlock()
 	if linkedSpecAuth != "Bearer test-token" {
 		t.Fatalf("linked spec Authorization = %q, want %q", linkedSpecAuth, "Bearer test-token")
 	}

internal/spec/discover.go

@@ -413,7 +413,10 @@ func discoverFromNetwork(ctx context.Context, cfg DiscoverConfig, loaders []Load
 			wg.Wait()
 			close(ch)
 		}()
-		return collectDiscoveryResults(ctx, cancel, ch, cfg.BaseURL)
+		spec, ttl, err := collectDiscoveryResults(ctx, cancel, ch, cfg.BaseURL)
+		cancel()
+		wg.Wait()
+		return spec, ttl, err
 	}
 
 	// Probe base URL: extract Link headers and try the body itself.
@@ -450,7 +453,10 @@ func discoverFromNetwork(ctx context.Context, cfg DiscoverConfig, loaders []Load
 		close(ch)
 	}()
 
-	return collectDiscoveryResults(ctx, cancel, ch, cfg.BaseURL)
+	spec, ttl, err := collectDiscoveryResults(ctx, cancel, ch, cfg.BaseURL)
+	cancel()
+	wg.Wait()
+	return spec, ttl, err
 }
 
 var wellKnownSpecPaths = []string{"/openapi.json", "/openapi.yaml"}

internal/spec/discover_test.go

@@ -884,6 +884,60 @@ func TestDiscover_ExplicitSpecURL(t *testing.T) {
 	}
 }
 
+func TestDiscoverWaitsForCanceledNetworkProbes(t *testing.T) {
+	specBody := `{"openapi":"3.1.0","info":{"title":"Direct","version":"1.0.0"},"paths":{}}`
+	slowEntered := make(chan struct{})
+	cancelSeen := make(chan struct{})
+	releaseSlow := make(chan struct{})
+	done := make(chan error, 1)
+
+	cfg := DiscoverConfig{
+		APIName: "testapi",
+		BaseURL: "https://api.example.com",
+		Fetch: func(ctx context.Context, rawURL string, tr http.RoundTripper) (*http.Response, error) {
+			switch rawURL {
+			case "https://api.example.com":
+				close(slowEntered)
+				<-ctx.Done()
+				close(cancelSeen)
+				<-releaseSlow
+				return nil, ctx.Err()
+			case "https://api.example.com/openapi.json":
+				<-slowEntered
+				return httpResponse(200, "application/json", specBody, nil), nil
+			default:
+				return httpResponse(404, "text/plain", "not found", nil), nil
+			}
+		},
+	}
+
+	go func() {
+		_, err := Discover(context.Background(), cfg, DefaultLoaders())
+		done <- err
+	}()
+
+	select {
+	case <-cancelSeen:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Discover did not cancel the slower probe")
+	}
+	select {
+	case err := <-done:
+		t.Fatalf("Discover returned before canceled probes exited: %v", err)
+	default:
+	}
+
+	close(releaseSlow)
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("Discover: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Discover did not return after canceled probe exited")
+	}
+}
+
 func TestDiscoverCleansCredentialURLMetadataInCache(t *testing.T) {
 	raw := `{"openapi":"3.1.0","info":{"title":"Direct","version":"1.0.0"},"paths":{"/items":{"get":{"operationId":"listItems","responses":{"200":{"description":"OK"}}}}}}`
 	tr := roundTripperFunc(func(r *http.Request) (*http.Response, error) {