fix(feature-flags): align agentless HTTP behavior

Author: leoromanovsky

products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/AgentlessConfigurationSource.java

@@ -1,16 +1,19 @@
 package com.datadog.featureflag;
 
 import static datadog.communication.http.OkHttpUtils.prepareRequest;
+import static datadog.communication.http.OkHttpUtils.sendWithRetries;
 import static datadog.trace.util.AgentThreadFactory.AgentThread.FEATURE_FLAG_CONFIGURATION_POLLER;
 import static datadog.trace.util.Strings.isBlank;
 
+import datadog.communication.http.HttpRetryPolicy;
 import datadog.communication.http.OkHttpUtils;
 import datadog.logging.RatelimitedLogger;
 import datadog.trace.api.Config;
 import datadog.trace.api.featureflag.FeatureFlaggingGateway;
 import datadog.trace.api.featureflag.ufc.v1.ServerConfiguration;
 import datadog.trace.util.AgentThreadFactory;
 import java.io.IOException;
+import java.io.InterruptedIOException;
 import java.net.HttpURLConnection;
 import java.util.HashMap;
 import java.util.Map;
@@ -50,8 +53,6 @@ final class AgentlessConfigurationSource implements ConfigurationSourceService {
   private final long pollIntervalMillis;
   private final UfcHttpClient client;
   private final ScheduledExecutorService executor;
-  private final RetrySleeper retrySleeper;
-  private final DoubleSupplier jitter;
   private final RatelimitedLogger ratelimitedLogger;
   private final Object lifecycleLock = new Object();
   private final AtomicBoolean polling = new AtomicBoolean();
@@ -71,11 +72,12 @@ private AgentlessConfigurationSource(final Config config, final HttpUrl endpoint
         new OkHttpUfcHttpClient(
             OkHttpUtils.buildHttpClient(
                 endpoint,
-                millis(config.getFeatureFlaggingConfigurationSourceRequestTimeoutSeconds()))),
+                millis(config.getFeatureFlaggingConfigurationSourceRequestTimeoutSeconds())),
+            millis(config.getFeatureFlaggingConfigurationSourcePollIntervalSeconds()),
+            TimeUnit.MILLISECONDS::sleep,
+            () -> ThreadLocalRandom.current().nextDouble(1 - RETRY_JITTER, 1 + RETRY_JITTER)),
         Executors.newSingleThreadScheduledExecutor(
             new AgentThreadFactory(FEATURE_FLAG_CONFIGURATION_POLLER)),
-        TimeUnit.MILLISECONDS::sleep,
-        () -> ThreadLocalRandom.current().nextDouble(1 - RETRY_JITTER, 1 + RETRY_JITTER),
         new RatelimitedLogger(LOGGER, MINUTES_BETWEEN_WARNINGS, TimeUnit.MINUTES));
   }
 
@@ -91,8 +93,6 @@ private AgentlessConfigurationSource(final Config config, final HttpUrl endpoint
         pollIntervalMillis,
         client,
         executor,
-        TimeUnit.MILLISECONDS::sleep,
-        () -> 1.0,
         new RatelimitedLogger(LOGGER, MINUTES_BETWEEN_WARNINGS, TimeUnit.MINUTES));
   }
 
@@ -102,35 +102,12 @@ private AgentlessConfigurationSource(final Config config, final HttpUrl endpoint
       final long pollIntervalMillis,
       final UfcHttpClient client,
       final ScheduledExecutorService executor,
-      final RetrySleeper retrySleeper,
-      final DoubleSupplier jitter) {
-    this(
-        endpoint,
-        config,
-        pollIntervalMillis,
-        client,
-        executor,
-        retrySleeper,
-        jitter,
-        new RatelimitedLogger(LOGGER, MINUTES_BETWEEN_WARNINGS, TimeUnit.MINUTES));
-  }
-
-  AgentlessConfigurationSource(
-      final HttpUrl endpoint,
-      final Config config,
-      final long pollIntervalMillis,
-      final UfcHttpClient client,
-      final ScheduledExecutorService executor,
-      final RetrySleeper retrySleeper,
-      final DoubleSupplier jitter,
       final RatelimitedLogger ratelimitedLogger) {
     this.endpoint = endpoint;
     this.config = config;
     this.pollIntervalMillis = pollIntervalMillis;
     this.client = client;
     this.executor = executor;
-    this.retrySleeper = retrySleeper;
-    this.jitter = jitter;
     this.ratelimitedLogger = ratelimitedLogger;
   }
 
@@ -184,55 +161,28 @@ private void pollOnceSafely() {
   }
 
   private boolean fetchAndApply() {
-    for (int attempt = 1; ; attempt++) {
-      try {
-        final UfcHttpResponse response = client.fetch(endpoint, config, etag);
-        if (closed) {
-          return false;
-        }
-        if (isRetryableStatus(response.status)) {
-          if (attempt < MAX_ATTEMPTS) {
-            if (!waitBeforeRetry(attempt)) {
-              return false;
-            }
-            continue;
-          }
-          ratelimitedLogger.warn(
-              "Feature Flagging agentless endpoint failed after {} attempts with HTTP {}",
-              MAX_ATTEMPTS,
-              response.status);
-          return false;
-        }
-        synchronized (lifecycleLock) {
-          return !closed && apply(response);
-        }
-      } catch (final IOException e) {
-        if (closed) {
-          return false;
-        }
-        if (attempt == MAX_ATTEMPTS) {
-          ratelimitedLogger.warn(
-              "Feature Flagging agentless endpoint request failed after {} attempts",
-              MAX_ATTEMPTS,
-              e);
-          return false;
-        }
-        if (!waitBeforeRetry(attempt)) {
-          return false;
-        }
-      }
-    }
-  }
-
-  private boolean waitBeforeRetry(final int attempt) {
-    if (closed) {
-      return false;
-    }
     try {
-      retrySleeper.sleep(retryDelayMillis(pollIntervalMillis, attempt, jitter.getAsDouble()));
-      return !closed;
-    } catch (final InterruptedException e) {
-      Thread.currentThread().interrupt();
+      final UfcHttpResponse response = client.fetch(endpoint, config, etag);
+      if (closed) {
+        return false;
+      }
+      if (isRetryableStatus(response.status)) {
+        ratelimitedLogger.warn(
+            "Feature Flagging agentless endpoint failed after {} attempts with HTTP {}",
+            MAX_ATTEMPTS,
+            response.status);
+        return false;
+      }
+      synchronized (lifecycleLock) {
+        return !closed && apply(response);
+      }
+    } catch (final IOException e) {
+      if (!closed) {
+        ratelimitedLogger.warn(
+            "Feature Flagging agentless endpoint request failed after {} attempts",
+            MAX_ATTEMPTS,
+            e);
+      }
       return false;
     }
   }
@@ -244,7 +194,7 @@ private boolean apply(final UfcHttpResponse response) {
     if (response.status == HttpURLConnection.HTTP_UNAUTHORIZED
         || response.status == HttpURLConnection.HTTP_FORBIDDEN) {
       ratelimitedLogger.warn(
-          "Feature Flagging agentless endpoint returned HTTP {}; verify DD_API_KEY is configured and valid",
+          "Feature Flagging agentless endpoint returned HTTP {}; verify endpoint authentication",
           response.status);
       return false;
     }
@@ -356,11 +306,30 @@ static final class UfcHttpResponse {
 
   static final class OkHttpUfcHttpClient implements UfcHttpClient {
     private final OkHttpClient httpClient;
-    private final AtomicReference<Call> activeCall = new AtomicReference<>();
+    private final long pollIntervalMillis;
+    private final RetrySleeper retrySleeper;
+    private final DoubleSupplier jitter;
+    private final AtomicBoolean fetching = new AtomicBoolean();
     private final AtomicBoolean cancelled = new AtomicBoolean();
+    private final AtomicReference<Call> activeCall = new AtomicReference<>();
 
     OkHttpUfcHttpClient(final OkHttpClient httpClient) {
+      this(
+          httpClient,
+          TimeUnit.SECONDS.toMillis(30),
+          TimeUnit.MILLISECONDS::sleep,
+          () -> ThreadLocalRandom.current().nextDouble(1 - RETRY_JITTER, 1 + RETRY_JITTER));
+    }
+
+    OkHttpUfcHttpClient(
+        final OkHttpClient httpClient,
+        final long pollIntervalMillis,
+        final RetrySleeper retrySleeper,
+        final DoubleSupplier jitter) {
       this.httpClient = httpClient;
+      this.pollIntervalMillis = pollIntervalMillis;
+      this.retrySleeper = retrySleeper;
+      this.jitter = jitter;
     }
 
     @Override
@@ -371,22 +340,50 @@ public UfcHttpResponse fetch(final HttpUrl endpoint, final Config config, final
         headers.put("If-None-Match", etag);
       }
       // Leave Accept-Encoding unset so OkHttp negotiates gzip and transparently decompresses it.
-      final Request request = prepareRequest(endpoint, headers, config, true).get().build();
-      final Call call = httpClient.newCall(request);
-      if (!activeCall.compareAndSet(null, call)) {
+      final Request request =
+          prepareRequest(endpoint, headers, config, isDatadogManagedEndpoint(endpoint, config))
+              .get()
+              .build();
+      if (!fetching.compareAndSet(false, true)) {
         throw new IllegalStateException("Feature Flagging HTTP request already in flight");
       }
       if (cancelled.get()) {
-        call.cancel();
+        fetching.set(false);
+        throw new InterruptedIOException("Feature Flagging HTTP client is closed");
       }
       try {
-        final Response response = call.execute();
-        try (ResponseBody responseBody = response.body()) {
-          final byte[] body = responseBody != null ? responseBody.bytes() : null;
-          return new UfcHttpResponse(response.code(), response.header("ETag"), body);
-        }
+        final HttpRetryPolicy.Factory retryPolicyFactory =
+            new HttpRetryPolicy.Factory(0, 0, 0) {
+              @Override
+              public HttpRetryPolicy create() {
+                return new AgentlessRetryPolicy(
+                    cancelled, pollIntervalMillis, retrySleeper, jitter);
+              }
+            };
+        final Call.Factory callFactory =
+            retryRequest -> {
+              final Call call = httpClient.newCall(retryRequest);
+              activeCall.set(call);
+              if (cancelled.get()) {
+                call.cancel();
+              }
+              return call;
+            };
+        return sendWithRetries(
+            callFactory,
+            retryPolicyFactory,
+            request,
+            response -> {
+              final int status = response.code();
+              final String responseEtag = response.header("ETag");
+              try (ResponseBody responseBody = response.body()) {
+                final byte[] body = responseBody != null ? responseBody.bytes() : null;
+                return new UfcHttpResponse(status, responseEtag, body);
+              }
+            });
       } finally {
-        activeCall.compareAndSet(call, null);
+        activeCall.set(null);
+        fetching.set(false);
       }
     }
 
@@ -398,5 +395,67 @@ public void cancel() {
         call.cancel();
       }
     }
+
+    private static boolean isDatadogManagedEndpoint(final HttpUrl endpoint, final Config config) {
+      return config.getFeatureFlaggingConfigurationSourceAgentlessBaseUrl() == null
+          && endpoint.isHttps()
+          && endpoint.host().equalsIgnoreCase("ufc-server.ff-cdn." + config.getSite());
+    }
+  }
+
+  static final class AgentlessRetryPolicy extends HttpRetryPolicy {
+    private final AtomicBoolean cancelled;
+    private final long pollIntervalMillis;
+    private final RetrySleeper retrySleeper;
+    private final DoubleSupplier jitter;
+    private int retriesLeft = MAX_ATTEMPTS - 1;
+    private int retryAttempt;
+
+    AgentlessRetryPolicy(
+        final AtomicBoolean cancelled,
+        final long pollIntervalMillis,
+        final RetrySleeper retrySleeper,
+        final DoubleSupplier jitter) {
+      super(0, 0, 0, false);
+      this.cancelled = cancelled;
+      this.pollIntervalMillis = pollIntervalMillis;
+      this.retrySleeper = retrySleeper;
+      this.jitter = jitter;
+    }
+
+    @Override
+    public boolean shouldRetry(final Exception exception) {
+      return exception instanceof IOException
+          && !Thread.currentThread().isInterrupted()
+          && reserveRetry();
+    }
+
+    @Override
+    public boolean shouldRetry(@Nullable final Response response) {
+      return response != null && isRetryableStatus(response.code()) && reserveRetry();
+    }
+
+    private boolean reserveRetry() {
+      if (cancelled.get() || retriesLeft == 0) {
+        return false;
+      }
+      retriesLeft--;
+      retryAttempt++;
+      return true;
+    }
+
+    @Override
+    public void backoff() throws IOException {
+      if (cancelled.get()) {
+        throw new InterruptedIOException("Feature Flagging HTTP client is closed");
+      }
+      try {
+        retrySleeper.sleep(
+            retryDelayMillis(pollIntervalMillis, retryAttempt, jitter.getAsDouble()));
+      } catch (final InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new InterruptedIOException("Feature Flagging retry interrupted");
+      }
+    }
   }
 }