[couchbase] Explicitly configure service quotas (and allow customization). (#4802)

This changeset explicitly configures the service memory quotas to their
documented limits by default. This ensures that if a service is enabled,
it does not consume more memory than absolutely needed. This is important
because some services (like analytics) try to grab memory available which
is good for production but of limited use in a testcontainers setting.

A new API is also added which allows to customize the memory quota for
each service should the test case require more than the minimum (which
for most services is 256MB anyways).

Author: daschl

modules/couchbase/src/main/java/org/testcontainers/couchbase/CouchbaseContainer.java

@@ -36,8 +36,10 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -100,6 +102,11 @@ public class CouchbaseContainer extends GenericContainer<CouchbaseContainer> {
         CouchbaseService.INDEX
     );
 
+    /**
+     * Holds the custom service quotas if configured by the user.
+     */
+    private final Map<CouchbaseService, Integer> customServiceQuotas = new HashMap<>();
+
     private final List<BucketDefinition> buckets = new ArrayList<>();
 
     private boolean isEnterprise = false;
@@ -158,6 +165,26 @@ public CouchbaseContainer withEnabledServices(final CouchbaseService... enabled)
         return this;
     }
 
+    /**
+     * Configures a custom memory quota for a given service.
+     *
+     * @param service the service to configure the quota for.
+     * @param quotaMb the memory quota in MB.
+     * @return this {@link CouchbaseContainer} for chaining purposes.
+     */
+    public CouchbaseContainer withServiceQuota(final CouchbaseService service, final int quotaMb) {
+        checkNotRunning();
+        if (!service.hasQuota()) {
+            throw new IllegalArgumentException("The provided service (" + service + ") has no quota to configure");
+        }
+        if (quotaMb < service.getMinimumQuotaMb()) {
+            throw new IllegalArgumentException("The custom quota (" + quotaMb + ") must not be smaller than the " +
+                "minimum quota for the service (" + service.getMinimumQuotaMb() + ")");
+        }
+        this.customServiceQuotas.put(service, quotaMb);
+        return this;
+    }
+
     /**
      * Enables the analytics service which is not enabled by default.
      *
@@ -262,6 +289,7 @@ protected void containerIsStarting(final InspectContainerResponse containerInfo)
         timePhase("initializeIsEnterprise", this::initializeIsEnterprise);
         timePhase("renameNode", this::renameNode);
         timePhase("initializeServices", this::initializeServices);
+        timePhase("setMemoryQuotas", this::setMemoryQuotas);
         timePhase("configureAdminUser", this::configureAdminUser);
         timePhase("configureExternalPorts", this::configureExternalPorts);
 
@@ -341,6 +369,35 @@ private void initializeServices() {
         checkSuccessfulResponse(response, "Could not enable couchbase services");
     }
 
+    /**
+     * Sets the memory quotas for each enabled service.
+     * <p>
+     * If there is no explicit custom quota defined, the default minimum quota will be used.
+     */
+    private void setMemoryQuotas() {
+        logger().debug("Custom service memory quotas: {}", customServiceQuotas);
+
+        final FormBody.Builder quotaBuilder = new FormBody.Builder();
+        for (CouchbaseService service : enabledServices) {
+            if (!service.hasQuota()) {
+                continue;
+            }
+
+            int quota = customServiceQuotas.getOrDefault(service, service.getMinimumQuotaMb());
+            if (CouchbaseService.KV.equals(service)) {
+                quotaBuilder.add("memoryQuota", Integer.toString(quota));
+            } else {
+                quotaBuilder.add(service.getIdentifier() + "MemoryQuota", Integer.toString(quota));
+            }
+        }
+
+        @Cleanup Response response = doHttpRequest(
+            MGMT_PORT, "/pools/default", "POST", quotaBuilder.build(), false
+        );
+
+        checkSuccessfulResponse(response, "Could not configure service memory quotas");
+    }
+
     /**
      * Configures the admin user on the couchbase node.
      * <p>

modules/couchbase/src/main/java/org/testcontainers/couchbase/CouchbaseService.java

@@ -21,38 +21,67 @@
  */
 public enum CouchbaseService {
 
-  /**
-   * Key-Value service.
-   */
-  KV("kv"),
-
-  /**
-   * Query (N1QL) service.
-   */
-  QUERY("n1ql"),
-
-  /**
-   * Search (FTS) service.
-   */
-  SEARCH("fts"),
-
-  /**
-   * Indexing service (needed if QUERY is also used!).
-   */
-  INDEX("index"),
-
-  /**
-   * Analytics service.
-   */
-  ANALYTICS("cbas");
-
-  private final String identifier;
-
-  CouchbaseService(String identifier) {
-    this.identifier = identifier;
-  }
-
-  String getIdentifier() {
-    return identifier;
-  }
+    /**
+     * Key-Value service.
+     */
+    KV("kv", 256),
+
+    /**
+     * Query (N1QL) service.
+     * <p>
+     * Note that the query service has no memory quota, so it is set to 0.
+     */
+    QUERY("n1ql", 0),
+
+    /**
+     * Search (FTS) service.
+     */
+    SEARCH("fts", 256),
+
+    /**
+     * Indexing service (needed if QUERY is also used!).
+     */
+    INDEX("index", 256),
+
+    /**
+     * Analytics service.
+     */
+    ANALYTICS("cbas", 1024);
+
+    private final String identifier;
+
+    private final int minimumQuotaMb;
+
+    CouchbaseService(final String identifier, final int minimumQuotaMb) {
+        this.identifier = identifier;
+        this.minimumQuotaMb = minimumQuotaMb;
+    }
+
+    /**
+     * Returns the internal service identifier.
+     *
+     * @return the internal service identifier.
+     */
+    String getIdentifier() {
+        return identifier;
+    }
+
+    /**
+     * Returns the minimum quota for the service in MB.
+     *
+     * @return the minimum quota in MB.
+     */
+    int getMinimumQuotaMb() {
+        return minimumQuotaMb;
+    }
+
+    /**
+     * Returns true if the service has a quota that needs to be applied.
+     *
+     * @return true if its quota needs to be applied.
+     */
+    boolean hasQuota() {
+        return minimumQuotaMb > 0;
+    }
+
 }