Implementation of Big Segments store support. (#15)

Author: gwhelanLD

build.gradle

@@ -38,7 +38,7 @@ allprojects {
 }
 
 ext.versions = [
-    "sdk": "5.0.0", // the *lowest* version we're compatible with
+    "sdk": "5.7.0", // the *lowest* version we're compatible with
     "dynamodb": "2.10.32",
     "slf4j": "1.7.21"
 ]

src/main/java/com/launchdarkly/sdk/server/integrations/DynamoDb.java

@@ -1,5 +1,8 @@
 package com.launchdarkly.sdk.server.integrations;
 
+import com.launchdarkly.sdk.server.interfaces.BigSegmentStoreFactory;
+import com.launchdarkly.sdk.server.interfaces.PersistentDataStoreFactory;
+
 /**
  * Integration between the LaunchDarkly SDK and DynamoDB.
  *
@@ -9,18 +12,39 @@ public abstract class DynamoDb {
   /**
    * Returns a builder object for creating a DynamoDB-backed data store.
    * <p>
-   * This object can be modified with {@link DynamoDbDataStoreBuilder} methods for any desired
-   * custom DynamoDB options. Then, pass it to
-   * {@link com.launchdarkly.sdk.server.Components#persistentDataStore(com.launchdarkly.sdk.server.interfaces.PersistentDataStoreFactory)}
-   * and set any desired caching options. Finally, pass the result to
-   * {@link com.launchdarkly.sdk.server.LDConfig.Builder#dataStore(com.launchdarkly.sdk.server.interfaces.DataStoreFactory)}.
-   * For example:
-   * 
+   * This can be used either for the main data store that holds feature flag data, or for the Big
+   * Segment store, or both. If you are using both, they do not have to have the same parameters.
+   * For instance, in this example the main data store uses a table called "table1" and the Big
+   * Segment store uses a table called "table2":
+   *
+   * <pre><code>
+   *   LDConfig config = new LDConfig.Builder()
+   *       .dataStore(
+   *           Components.persistentDataStore(
+   *               DynamoDb.dataStore("table1")
+   *           )
+   *       )
+   *       .bigSegments(
+   *           Components.bigSegments(
+   *               DynamoDb.dataStore("table2")
+   *           )
+   *       )
+   *       .build();
+   * </code></pre>
+   *
+   * Note that the builder is passed to one of two methods,
+   * {@link com.launchdarkly.sdk.server.Components#persistentDataStore(PersistentDataStoreFactory)} or
+   * {@link com.launchdarkly.sdk.server.Components#bigSegments(BigSegmentStoreFactory)}, depending on
+   * the context in which it is being used. This is because each of those contexts has its own
+   * additional configuration options that are unrelated to the DynamoDb options. For instance, the
+   * {@link com.launchdarkly.sdk.server.Components#persistentDataStore(PersistentDataStoreFactory)}
+   * builder has options for caching:
+   *
    * <pre><code>
    *     LDConfig config = new LDConfig.Builder()
    *         .dataStore(
    *             Components.persistentDataStore(
-   *                 DynamoDb.dataStore("my-table-name")
+   *                 DynamoDb.dataStore("table1")
    *             ).cacheSeconds(15)
    *         )
    *         .build();

src/main/java/com/launchdarkly/sdk/server/integrations/DynamoDbBigSegmentStoreImpl.java

@@ -0,0 +1,59 @@
+package com.launchdarkly.sdk.server.integrations;
+
+import static com.launchdarkly.sdk.server.interfaces.BigSegmentStoreTypes.createMembershipFromSegmentRefs;
+
+import com.launchdarkly.sdk.server.interfaces.BigSegmentStore;
+import com.launchdarkly.sdk.server.interfaces.BigSegmentStoreTypes;
+
+import java.util.List;
+
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
+
+public class DynamoDbBigSegmentStoreImpl extends DynamoDbStoreImplBase implements BigSegmentStore {
+  private final static String MEMBERSHIP_KEY = "big_segments_user";
+  private final static String INCLUDED_ATTR = "included";
+  private final static String EXCLUDED_ATTR = "excluded";
+
+  private final static String METADATA_KEY = "big_segments_metadata";
+  private final static String SYNC_TIME_ATTR = "synchronizedOn";
+
+  DynamoDbBigSegmentStoreImpl(DynamoDbClient client, boolean wasExistingClient, String tableName, String prefix) {
+    super(client, wasExistingClient, tableName, prefix);
+  }
+
+  @Override
+  public BigSegmentStoreTypes.Membership getMembership(String userHash) {
+    String namespaceKey = prefixedNamespace(MEMBERSHIP_KEY);
+    GetItemResponse response = getItemByKeys(namespaceKey, userHash);
+    if (response == null || response.item() == null || response.item().isEmpty()) {
+      return null;
+    }
+    List<String> includedRefs = stringListFromAttrValue(response.item().get(INCLUDED_ATTR));
+    List<String> excludedRefs = stringListFromAttrValue(response.item().get(EXCLUDED_ATTR));
+    return createMembershipFromSegmentRefs(includedRefs, excludedRefs);
+  }
+
+  private static List<String> stringListFromAttrValue(AttributeValue attrValue) {
+    return attrValue == null ? null : attrValue.ss();
+  }
+
+  @Override
+  public BigSegmentStoreTypes.StoreMetadata getMetadata() {
+    String key = prefixedNamespace(METADATA_KEY);
+    GetItemResponse response = getItemByKeys(key, key);
+    if (response == null || response.item() == null) {
+      return null;
+    }
+    AttributeValue syncTimeValue = response.item().get(SYNC_TIME_ATTR);
+    if (syncTimeValue == null) {
+      return null;
+    }
+    String syncTimeString = syncTimeValue.n();
+    if (syncTimeString == null) {
+      return null;
+    }
+    return new BigSegmentStoreTypes.StoreMetadata(Long.parseLong(syncTimeString));
+  }
+}

src/main/java/com/launchdarkly/sdk/server/integrations/DynamoDbDataStoreBuilder.java

@@ -1,8 +1,9 @@
 package com.launchdarkly.sdk.server.integrations;
 
 import com.launchdarkly.sdk.LDValue;
-import com.launchdarkly.sdk.server.LDConfig;
 import com.launchdarkly.sdk.server.interfaces.BasicConfiguration;
+import com.launchdarkly.sdk.server.interfaces.BigSegmentStore;
+import com.launchdarkly.sdk.server.interfaces.BigSegmentStoreFactory;
 import com.launchdarkly.sdk.server.interfaces.ClientContext;
 import com.launchdarkly.sdk.server.interfaces.DiagnosticDescription;
 import com.launchdarkly.sdk.server.interfaces.PersistentDataStore;
@@ -19,22 +20,17 @@
 /**
  * Builder/factory class for the DynamoDB data store.
  * <p>
- * Obtain an instance of this class by calling {@link com.launchdarkly.sdk.server.integrations.DynamoDb#dataStore(String)}.
- * After calling its methods to specify any desired custom settings, wrap it in a
- * {@link com.launchdarkly.sdk.server.integrations.PersistentDataStoreBuilder}
- * by calling {@code Components.persistentDataStore()}, then pass the result into the SDK configuration with
- * {@link com.launchdarkly.sdk.server.LDConfig.Builder#dataStore(com.launchdarkly.sdk.server.interfaces.DataStoreFactory)}.
- * You do not need to call {@link #createPersistentDataStore(ClientContext)} yourself to build the actual data store; that
- * will be done by the SDK.
+ * Examples of configuring the SDK with this builder are described in the documentation for
+ * {@link DynamoDb#dataStore(String)}, which returns an instance of this class.
  * <p>
  * The AWS SDK provides many configuration options for a DynamoDB client. This class has
  * corresponding methods for some of the most commonly used ones. If you need more sophisticated
  * control over the DynamoDB client, you can construct one of your own and pass it in with the
  * {@link #existingClient(DynamoDbClient)} method.
- * 
+ *
  * @since 2.1.0
  */
-public final class DynamoDbDataStoreBuilder implements PersistentDataStoreFactory, DiagnosticDescription {
+public final class DynamoDbDataStoreBuilder implements PersistentDataStoreFactory, BigSegmentStoreFactory, DiagnosticDescription {
   private final String tableName;
 
   private String prefix;
@@ -111,7 +107,7 @@ public DynamoDbDataStoreBuilder prefix(String prefix) {
    * Specifies an existing, already-configured DynamoDB client instance that the data store
    * should use rather than creating one of its own. If you specify an existing client, then the
    * other builder methods for configuring DynamoDB are ignored.
-   *  
+   *
    * @param existingClient an existing DynamoDB client instance
    * @return the builder
    */
@@ -127,7 +123,17 @@ public DynamoDbDataStoreBuilder existingClient(DynamoDbClient existingClient) {
   @Override
   public PersistentDataStore createPersistentDataStore(ClientContext context) {  
     DynamoDbClient client = (existingClient != null) ? existingClient : clientBuilder.build();
-    return new DynamoDbDataStoreImpl(client, tableName, prefix);
+    return new DynamoDbDataStoreImpl(client, existingClient != null, tableName, prefix);
+  }
+
+  /**
+   * Called internally by the SDK to create the actual Big Segment store instance.
+   * @return the Big Segment store configured by this builder
+   */
+  @Override
+  public BigSegmentStore createBigSegmentStore(ClientContext context) {
+    DynamoDbClient client = (existingClient != null) ? existingClient : clientBuilder.build();
+    return new DynamoDbBigSegmentStoreImpl(client, existingClient != null, tableName, prefix);
   }
 
   @Override

src/main/java/com/launchdarkly/sdk/server/integrations/DynamoDbDataStoreImpl.java

@@ -9,7 +9,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.IOException;
 import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -59,29 +58,17 @@
  * stored as a single item, this mechanism will not work for extremely large flags or segments.
  * </ul>
  */
-final class DynamoDbDataStoreImpl implements PersistentDataStore {
+final class DynamoDbDataStoreImpl extends DynamoDbStoreImplBase implements PersistentDataStore {
   private static final Logger logger = LoggerFactory.getLogger("com.launchdarkly.sdk.server.LDClient.DataStore.DynamoDB");
 
-  static final String partitionKey = "namespace";
-  static final String sortKey = "key";
   private static final String versionAttribute = "version";
   private static final String itemJsonAttribute = "item";
   private static final String deletedItemPlaceholder = "null"; // DynamoDB doesn't allow empty strings
-  private final DynamoDbClient client;
-  private final String tableName;
-  private final String prefix;
-  
+
   private Runnable updateHook;
 
-  DynamoDbDataStoreImpl(DynamoDbClient client, String tableName, String prefix) {
-    this.client = client;
-    this.tableName = tableName;
-    this.prefix = "".equals(prefix) ? null : prefix;
-  }
-  
-  @Override
-  public void close() throws IOException {
-    client.close();
+  DynamoDbDataStoreImpl(DynamoDbClient client, boolean wasExistingClient, String tableName, String prefix) {
+    super(client, wasExistingClient, tableName, prefix);
   }
 
   @Override
@@ -95,10 +82,8 @@ public KeyedItems<SerializedItemDescriptor> getAll(DataKind kind) {
     List<Map.Entry<String, SerializedItemDescriptor>> itemsOut = new ArrayList<>();
     for (QueryResponse resp: client.queryPaginator(makeQueryForKind(kind).build())) {
       for (Map<String, AttributeValue> item: resp.items()) {
-        AttributeValue keyAttr = item.get(sortKey);
-        if (keyAttr == null || keyAttr.s() == null) {
-          
-        } else {
+        AttributeValue keyAttr = item.get(SORT_KEY);
+        if (keyAttr != null && keyAttr.s() != null) {
           SerializedItemDescriptor itemOut = unmarshalItem(kind, item);
           if (itemOut != null) {
             itemsOut.add(new AbstractMap.SimpleEntry<>(keyAttr.s(), itemOut));
@@ -135,18 +120,17 @@ public void init(FullDataSet<SerializedItemDescriptor> allData) {
     // Now delete any previously existing items whose keys were not in the current data
     for (Map.Entry<String, String> combinedKey: unusedOldKeys) {
       if (!combinedKey.getKey().equals(initedKey())) {
-        Map<String, AttributeValue> keys = mapOf(
-            partitionKey, AttributeValue.builder().s(combinedKey.getKey()).build(),
-            sortKey, AttributeValue.builder().s(combinedKey.getValue()).build());
-        requests.add(WriteRequest.builder().deleteRequest(builder -> builder.key(keys)).build());
+        requests.add(WriteRequest.builder()
+            .deleteRequest(builder ->
+                builder.key(makeKeysMap(combinedKey.getKey(), combinedKey.getValue())))
+            .build());
       }
     }
 
     // Now set the special key that we check in initializedInternal()
-    Map<String, AttributeValue> initedItem = mapOf(
-        partitionKey, AttributeValue.builder().s(initedKey()).build(),
-        sortKey, AttributeValue.builder().s(initedKey()).build());
-    requests.add(WriteRequest.builder().putRequest(builder -> builder.item(initedItem)).build());
+    requests.add(WriteRequest.builder()
+        .putRequest(builder -> builder.item(makeKeysMap(initedKey(), initedKey())))
+        .build());
 
     batchWriteRequests(client, tableName, requests);
 
@@ -166,8 +150,8 @@ public boolean upsert(DataKind kind, String key, SerializedItemDescriptor newIte
           .item(encodedItem)
           .conditionExpression("attribute_not_exists(#namespace) or attribute_not_exists(#key) or :version > #version")
           .expressionAttributeNames(mapOf(
-              "#namespace", partitionKey,
-              "#key", sortKey,
+              "#namespace", PARTITION_KEY,
+              "#key", SORT_KEY,
               "#version", versionAttribute))
           .expressionAttributeValues(mapOf(
               ":version", AttributeValue.builder().n(String.valueOf(newItem.getVersion())).build()))
@@ -199,11 +183,7 @@ public boolean isStoreAvailable() {
   public void setUpdateHook(Runnable updateHook) {
     this.updateHook = updateHook;
   }
-  
-  private String prefixedNamespace(String base) {
-    return prefix == null ? base : (prefix + ":" + base);
-  }
-  
+
   private String namespaceForKind(DataKind kind) {
     return prefixedNamespace(kind.getName());
   }
@@ -214,7 +194,7 @@ private String initedKey() {
 
   private QueryRequest.Builder makeQueryForKind(DataKind kind) {
     Map<String, Condition> keyConditions = mapOf(
-        partitionKey,
+        PARTITION_KEY,
         Condition.builder()
           .comparisonOperator(ComparisonOperator.EQ)
           .attributeValueList(AttributeValue.builder().s(namespaceForKind(kind)).build())
@@ -226,31 +206,20 @@ private QueryRequest.Builder makeQueryForKind(DataKind kind) {
         .keyConditions(keyConditions);
   }
 
-  private GetItemResponse getItemByKeys(String namespace, String key) {
-    Map<String, AttributeValue> keyMap = mapOf(
-        partitionKey, AttributeValue.builder().s(namespace).build(),
-        sortKey, AttributeValue.builder().s(key).build()
-    );
-    return client.getItem(builder -> builder.tableName(tableName)
-        .consistentRead(true)
-        .key(keyMap)
-    );
-  }
-  
   private Set<Map.Entry<String, String>> readExistingKeys(FullDataSet<?> kindsFromThisDataSet) {
     Set<Map.Entry<String, String>> keys = new HashSet<>();
     for (Map.Entry<DataKind, ?> e: kindsFromThisDataSet.getData()) {
       DataKind kind = e.getKey();
       QueryRequest req = makeQueryForKind(kind)
         .projectionExpression("#namespace, #key")
         .expressionAttributeNames(mapOf(
-            "#namespace", partitionKey, "#key", sortKey))
+            "#namespace", PARTITION_KEY, "#key", SORT_KEY))
         .build();
       QueryIterable queryResults = client.queryPaginator(req);
       for (QueryResponse resp: queryResults) {
         for (Map<String, AttributeValue> item: resp.items()) {
-          String namespace = item.get(partitionKey).s();
-          String key = item.get(sortKey).s();
+          String namespace = item.get(PARTITION_KEY).s();
+          String key = item.get(SORT_KEY).s();
           keys.add(new AbstractMap.SimpleEntry<>(namespace, key));
         }
       }
@@ -261,8 +230,8 @@ private Set<Map.Entry<String, String>> readExistingKeys(FullDataSet<?> kindsFrom
   private Map<String, AttributeValue> marshalItem(DataKind kind, String key, SerializedItemDescriptor item) {
     String json = item.isDeleted() ? deletedItemPlaceholder : item.getSerializedItem();
     return mapOf(
-        partitionKey, AttributeValue.builder().s(namespaceForKind(kind)).build(),
-        sortKey, AttributeValue.builder().s(key).build(),
+        PARTITION_KEY, AttributeValue.builder().s(namespaceForKind(kind)).build(),
+        SORT_KEY, AttributeValue.builder().s(key).build(),
         versionAttribute, AttributeValue.builder().n(String.valueOf(item.getVersion())).build(),
         itemJsonAttribute, AttributeValue.builder().s(json).build()
     );