Cache flags DB

Up until now every lookup of a flag on ZooKeeperFlagSource would hit ZooKeeper.
Flags are ideal for caching: Changes seldom, little data, clients should handle
short-lived inconsistencies.

This PR will make the backing FlagsDbImpl cache the /flags/v1 ZK directory and
completes the optimization of ConfigServerFlagSource.

2 files changed, 20 insertions, 9 deletions

diff --git a/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImpl.java b/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImpl.java
index 7b0a2f632cc..c313a47ad13 100644
--- a/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImpl.java
+++ b/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImpl.java
@@ -7,10 +7,15 @@ import com.yahoo.vespa.configserver.flags.FlagsDb;
 import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.flags.FlagId;
 import com.yahoo.vespa.flags.json.FlagData;
+import org.apache.curator.framework.recipes.cache.ChildData;
 
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 /**
  * @author hakonhall
@@ -19,15 +24,21 @@ public class FlagsDbImpl implements FlagsDb {
     private static final Path ROOT_PATH = Path.fromString("/flags/v1");
 
     private final Curator curator;
+    private final Curator.DirectoryCache cache;
 
     @Inject
     public FlagsDbImpl(Curator curator) {
         this.curator = curator;
+        curator.create(ROOT_PATH);
+        ExecutorService executorService = Executors.newFixedThreadPool(1);
+        this.cache = curator.createDirectoryCache(ROOT_PATH.toString(), true, false, executorService);
     }
 
     @Override
     public Optional<FlagData> getValue(FlagId flagId) {
-        return curator.getData(getZkPathFor(flagId)).map(FlagData::deserializeUtf8Json);
+        return Optional.ofNullable(cache.getCurrentData(getZkPathFor(flagId)))
+                .map(ChildData::getData)
+                .map(FlagData::deserializeUtf8Json);
     }
 
     @Override
@@ -37,12 +48,11 @@ public class FlagsDbImpl implements FlagsDb {
 
     @Override
     public Map<FlagId, FlagData> getAllFlags() {
-        Map<FlagId, FlagData> flags = new HashMap<>();
-        for (String flagId : curator.getChildren(ROOT_PATH)) {
-            FlagId id = new FlagId(flagId);
-            getValue(id).ifPresent(data -> flags.put(id, data));
-        }
-        return flags;
+        List<ChildData> dataList = cache.getCurrentData();
+        return dataList.stream()
+                .map(ChildData::getData)
+                .map(FlagData::deserializeUtf8Json)
+                .collect(Collectors.toMap(FlagData::id, Function.identity()));
     }
 
     @Override
diff --git a/configserver-flags/src/test/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImplTest.java b/configserver-flags/src/test/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImplTest.java
index 0523090c18b..ecc9bacb081 100644
--- a/configserver-flags/src/test/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImplTest.java
+++ b/configserver-flags/src/test/java/com/yahoo/vespa/configserver/flags/db/FlagsDbImplTest.java
@@ -43,7 +43,8 @@ public class FlagsDbImplTest {
                 dataCopy.get().serializeToJson());
 
         FlagId flagId2 = new FlagId("id2");
-        db.setValue(flagId2, data);
+        FlagData data2 = new FlagData(flagId2, new FetchVector().with(FetchVector.Dimension.ZONE_ID, "zone-a"), rule1);
+        db.setValue(flagId2, data2);
         Map<FlagId, FlagData> flags = db.getAllFlags();
         assertThat(flags.size(), equalTo(2));
         assertThat(flags.get(flagId), notNullValue());