Retry on failed updates

Signed-off-by: Byron Ruth <[email protected]>

pkg/drivers/nats/new.go

@@ -233,17 +233,26 @@ func getOrCreateBucket(ctx context.Context, js jetstream.JetStream, config *Conf
 }
 
 func disableDirectGets(ctx context.Context, js jetstream.JetStream, config *Config) error {
-	str, err := js.Stream(ctx, fmt.Sprintf("KV_%s", config.bucket))
-	if err != nil {
-		return fmt.Errorf("failed to get stream info: %w", err)
-	}
-	scfg := str.CachedInfo().Config
-	scfg.AllowDirect = false
+	for {
+		str, err := js.Stream(ctx, fmt.Sprintf("KV_%s", config.bucket))
+		if errors.Is(err, context.DeadlineExceeded) {
+			continue
+		}
+		if err != nil {
+			return fmt.Errorf("failed to get stream info: %w", err)
+		}
 
-	_, err = js.UpdateStream(ctx, scfg)
-	if err != nil {
-		return fmt.Errorf("failed to update stream config: %w", err)
-	}
+		scfg := str.CachedInfo().Config
+		scfg.AllowDirect = false
+
+		_, err = js.UpdateStream(ctx, scfg)
+		if errors.Is(err, context.DeadlineExceeded) {
+			continue
+		}
+		if err != nil {
+			return fmt.Errorf("failed to update stream config: %w", err)
+		}
 
-	return nil
+		return nil
+	}
 }