server: add a snapshot collecting DiscoveryServerCallbacks (#73)

Adds a callback that watches the streams for requests and stream closes
and cleans up unreferenced snapshots after some time period.

Signed-off-by: Snow Pettersen <[email protected]>

Author: snowp

cache/src/main/java/io/envoyproxy/controlplane/cache/Snapshot.java

@@ -15,6 +15,7 @@
 import envoy.api.v2.Lds.Listener;
 import envoy.api.v2.Rds.RouteConfiguration;
 import envoy.api.v2.auth.Cert.Secret;
+import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 
@@ -83,6 +84,16 @@ public static Snapshot create(
         SnapshotResources.create(secrets, secretsVersion));
   }
 
+  /**
+   * Creates an empty snapshot with the given version.
+   *
+   * @param version the version of the snapshot resources
+   */
+  public static Snapshot createEmpty(String version) {
+    return create(Collections.emptySet(), Collections.emptySet(),
+        Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), version);
+  }
+
   /**
    * Returns all cluster items in the CDS payload.
    */

server/src/main/java/io/envoyproxy/controlplane/server/callback/SnapshotCollectingCallback.java

@@ -0,0 +1,137 @@
+package io.envoyproxy.controlplane.server.callback;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import envoy.api.v2.Discovery;
+import io.envoyproxy.controlplane.cache.NodeGroup;
+import io.envoyproxy.controlplane.cache.Snapshot;
+import io.envoyproxy.controlplane.cache.SnapshotCache;
+import io.envoyproxy.controlplane.server.DiscoveryServerCallbacks;
+import java.time.Clock;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+/**
+ * Callback that keeps track of the number of streams associated with each node group and periodically clears
+ * out {@link Snapshot}s from the cache that are no longer referenced by any streams.
+ *
+ * <p>Works by monitoring the stream to determine what group they belong to and keeps a running count as well
+ * as when a request is seen that targets a given node group.
+ *
+ * <p>Every {@code collectionIntervalMillis} milliseconds a cleanup job runs which looks for snapshots with no
+ * active streams that haven't been updated within the configured time frame. Checking the time since last update
+ * is done to prevent snapshots from being prematurely removed from the cache. It ensures that a group must have
+ * no active streams for {@code collectAfterMillis} milliseconds before being collected.
+ *
+ * <p>To be notified of snapshots that are removed, a set of callbacks may be provided which will be triggered
+ * whenever a snapshot is removed from the cache. Any other callback which maintains state about the snapshots
+ * that is cleaned up by one of these callbacks should be run *after* this callback. This helps ensure that
+ * if state is cleaned up while a request in inbound, the request will be blocked by the lock in this callback
+ * until collection finishes and the subsequent callbacks will see the new request come in after collection. If the
+ * order is reversed, another callback might have seen the new request but the refcount here hasn't been incremented,
+ * causing it to get cleaned up and wipe the state of the other callback even though we now have an active stream
+ * for that group.
+ */
+public class SnapshotCollectingCallback<T> implements DiscoveryServerCallbacks {
+  private static class SnapshotState {
+    int streamCount;
+    Instant lastSeen;
+  }
+
+  private final SnapshotCache<T> snapshotCache;
+  private final NodeGroup<T> nodeGroup;
+  private final Clock clock;
+  private final Set<Consumer<T>> collectorCallbacks;
+  private final long collectAfterMillis;
+  private final Map<T, SnapshotState> snapshotStates = new ConcurrentHashMap<>();
+  private final Map<Long, T> groupByStream = new ConcurrentHashMap<>();
+
+  /**
+   * Creates the callback.
+   *
+   * @param snapshotCache the cache to evict snapshots from
+   * @param nodeGroup the node group used to map requests to groups
+   * @param clock system clock
+   * @param collectorCallbacks the callbacks to invoke when snapshot is collected
+   * @param collectAfterMillis how long a snapshot must be referenced for before being collected
+   * @param collectionIntervalMillis how often the collection background action should run
+   */
+  public SnapshotCollectingCallback(SnapshotCache<T> snapshotCache,
+      NodeGroup<T> nodeGroup, Clock clock, Set<Consumer<T>> collectorCallbacks,
+      long collectAfterMillis, long collectionIntervalMillis) {
+    this.snapshotCache = snapshotCache;
+    this.nodeGroup = nodeGroup;
+    this.clock = clock;
+    this.collectorCallbacks = collectorCallbacks;
+    this.collectAfterMillis = collectAfterMillis;
+    ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor(
+        new ThreadFactoryBuilder().setNameFormat("snapshot-gc-%d").build());
+    executorService.scheduleAtFixedRate(() -> deleteUnreferenced(clock), collectionIntervalMillis,
+        collectionIntervalMillis, TimeUnit.MILLISECONDS);
+  }
+
+  @Override
+  public synchronized void onStreamRequest(long streamId, Discovery.DiscoveryRequest request) {
+    T groupIdentifier = nodeGroup.hash(request.getNode());
+
+    SnapshotState snapshotState =
+        this.snapshotStates.computeIfAbsent(groupIdentifier, x -> new SnapshotState());
+    snapshotState.lastSeen = clock.instant();
+
+    if (groupByStream.put(streamId, groupIdentifier) == null) {
+      snapshotState.streamCount++;
+    }
+  }
+
+  @Override public void onStreamClose(long streamId, String typeUrl) {
+    onStreamCloseHelper(streamId);
+  }
+
+  @Override public void onStreamCloseWithError(long streamId, String typeUrl, Throwable error) {
+    onStreamCloseHelper(streamId);
+  }
+
+  @VisibleForTesting
+  synchronized void deleteUnreferenced(Clock clock) {
+    // Keep track of snapshots to delete to avoid CME.
+    Set<T> toDelete = new LinkedHashSet<>();
+
+    for (Map.Entry<T, SnapshotState> entry : snapshotStates.entrySet()) {
+      if (entry.getValue().streamCount == 0 && entry.getValue().lastSeen.isBefore(
+          clock.instant().minus(collectAfterMillis, ChronoUnit.MILLIS))) {
+
+        // clearSnapshot will do nothing and return false if there are any pending watches - this
+        // ensures that we don't actually remove a snapshot that's in use.
+        T groupIdentifier = entry.getKey();
+        if (snapshotCache.clearSnapshot(groupIdentifier)) {
+          toDelete.add(groupIdentifier);
+        }
+      }
+    }
+
+    toDelete.forEach(group -> {
+      snapshotStates.remove(group);
+      collectorCallbacks.forEach(cb -> cb.accept(group));
+    });
+  }
+
+  private synchronized void onStreamCloseHelper(long streamId) {
+    T removed = groupByStream.remove(streamId);
+    if (removed == null) {
+      // This will happen if the stream closed before we received the first request.
+      return;
+    }
+
+    SnapshotState snapshotState = snapshotStates.get(removed);
+    snapshotState.streamCount--;
+    snapshotState.lastSeen = clock.instant();
+  }
+}

server/src/test/java/io/envoyproxy/controlplane/server/callback/SnapshotCollectingCallbackTest.java

@@ -0,0 +1,92 @@
+package io.envoyproxy.controlplane.server.callback;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.google.common.collect.ImmutableSet;
+import envoy.api.v2.Discovery;
+import io.envoyproxy.controlplane.cache.NodeGroup;
+import io.envoyproxy.controlplane.cache.SimpleCache;
+import io.envoyproxy.controlplane.cache.Snapshot;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import org.junit.Before;
+import org.junit.Test;
+
+public class SnapshotCollectingCallbackTest {
+
+  private static final Clock CLOCK = Clock.fixed(Instant.now(), ZoneId.systemDefault());
+  private static final NodeGroup<String> NODE_GROUP = node -> "group";
+  private final ArrayList<String> collectedGroups = new ArrayList<>();
+  private SnapshotCollectingCallback<String> callback;
+  private SimpleCache<String> cache;
+
+  @Before
+  public void setUp() {
+    collectedGroups.clear();
+    cache = new SimpleCache<>(NODE_GROUP);
+    cache.setSnapshot("group", Snapshot.createEmpty(""));
+    callback = new SnapshotCollectingCallback<>(cache, NODE_GROUP, CLOCK,
+        Collections.singleton(collectedGroups::add), 3, 100);
+  }
+
+  @Test
+  public void testSingleSnapshot() {
+    callback.onStreamRequest(0, Discovery.DiscoveryRequest.getDefaultInstance());
+    callback.onStreamRequest(1, Discovery.DiscoveryRequest.getDefaultInstance());
+
+    // We have 2 references to the snapshot, this should do nothing.
+    callback.deleteUnreferenced(Clock.offset(CLOCK, Duration.ofMillis(5)));
+    assertThat(collectedGroups).isEmpty();
+
+    callback.onStreamClose(0, "");
+
+    // We have 1 reference to the snapshot, this should do nothing.
+    callback.deleteUnreferenced(Clock.offset(CLOCK, Duration.ofMillis(5)));
+    assertThat(collectedGroups).isEmpty();
+
+    callback.onStreamCloseWithError(1, "", new RuntimeException());
+
+    // We have 0 references to the snapshot, but 1 < 3 so it's too early to collect the snapshot.
+    callback.deleteUnreferenced(Clock.offset(CLOCK, Duration.ofMillis(1)));
+    assertThat(collectedGroups).isEmpty();
+
+    // We have 0 references to the snapshot, and 5 > 3 so we clear out the snapshot.
+    callback.deleteUnreferenced(Clock.offset(CLOCK, Duration.ofMillis(5)));
+    assertThat(collectedGroups).containsExactly("group");
+  }
+
+  @Test
+  public void testAsyncCollection() throws InterruptedException {
+    CountDownLatch snapshotCollectedLatch = new CountDownLatch(1);
+    CountDownLatch deleteUnreferencedLatch = new CountDownLatch(1);
+
+    // Create a cache with 0 expiry delay, which means the snapshot should get collected immediately.
+    callback = new SnapshotCollectingCallback<String>(cache, NODE_GROUP, CLOCK,
+        ImmutableSet.of(collectedGroups::add, group -> snapshotCollectedLatch.countDown()), -3, 1) {
+      @Override synchronized void deleteUnreferenced(Clock clock) {
+        super.deleteUnreferenced(clock);
+        deleteUnreferencedLatch.countDown();;
+      }
+    };
+
+    callback.onStreamRequest(0, Discovery.DiscoveryRequest.getDefaultInstance());
+    assertThat(deleteUnreferencedLatch.await(100, TimeUnit.MILLISECONDS)).isTrue();
+    assertThat(collectedGroups).isEmpty();
+
+    callback.onStreamClose(0, "");
+    assertThat(snapshotCollectedLatch.await(100,TimeUnit.MILLISECONDS)).isTrue();
+    assertThat(collectedGroups).containsExactly("group");
+  }
+
+  @Test
+  public void testCloseBeforeRequest() {
+    callback.onStreamClose(0, "");
+    assertThat(collectedGroups).isEmpty();
+  }
+}