OAK-6223: Expose socket keep-alive option

git-svn-id: https://svn.apache.org/repos/asf/jackrabbit/oak/trunk@1795330 13f79535-47bb-0310-9956-ffa450edef68

oak-core/src/main/java/org/apache/jackrabbit/oak/plugins/document/DocumentMK.java

@@ -47,6 +47,7 @@
 import com.google.common.collect.Sets;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.mongodb.DB;
+import com.mongodb.MongoClientOptions;
 import com.mongodb.ReadConcernLevel;
 import org.apache.jackrabbit.oak.api.CommitFailedException;
 import org.apache.jackrabbit.oak.api.PropertyState;
@@ -556,6 +557,7 @@ public static class Builder {
         private DocumentNodeStore nodeStore;
         private DocumentStore documentStore;
         private String mongoUri;
+        private boolean socketKeepAlive;
         private MongoStatus mongoStatus;
         private DiffCache diffCache;
         private BlobStore blobStore;
@@ -623,7 +625,9 @@ public Builder setMongoDB(@Nonnull String uri,
                 throws UnknownHostException {
             this.mongoUri = uri;
 
-            DB db = new MongoConnection(uri).getDB(name);
+            MongoClientOptions.Builder options = MongoConnection.getDefaultBuilder();
+            options.socketKeepAlive(socketKeepAlive);
+            DB db = new MongoConnection(uri, options).getDB(name);
             MongoStatus status = new MongoStatus(db);
             if (!MongoConnection.hasWriteConcern(uri)) {
                 db.setWriteConcern(MongoConnection.getDefaultWriteConcern(db));
@@ -677,6 +681,18 @@ private Builder setMongoDB(@Nonnull DB db,
             return this;
         }
 
+        /**
+         * Enables the socket keep-alive option for MongoDB. The default is
+         * disabled.
+         *
+         * @param enable whether to enable it.
+         * @return this
+         */
+        public Builder setSocketKeepAlive(boolean enable) {
+            this.socketKeepAlive = enable;
+            return this;
+        }
+
         private void setBlobStore(GarbageCollectableBlobStore s) {
             configureBlobStore(s);
             PersistentCache p = getPersistentCache();

oak-core/src/main/java/org/apache/jackrabbit/oak/plugins/document/DocumentNodeStoreService.java

@@ -129,6 +129,7 @@ public class DocumentNodeStoreService {
     private static final int DEFAULT_CACHE = (int) (DEFAULT_MEMORY_CACHE_SIZE / MB);
     private static final int DEFAULT_BLOB_CACHE_SIZE = 16;
     private static final String DEFAULT_DB = "oak";
+    private static final boolean DEFAULT_SO_KEEP_ALIVE = false;
     private static final String DEFAULT_PERSISTENT_CACHE = "cache,binary=0";
     private static final String DEFAULT_JOURNAL_CACHE = "diff-cache";
     private static final String PREFIX = "oak.documentstore.";
@@ -145,6 +146,11 @@ public class DocumentNodeStoreService {
      */
     private static final String FWK_PROP_DB = "oak.mongo.db";
 
+    /**
+     * Name of framework property to configure socket keep-alive for MongoDB
+     */
+    private static final String FWK_PROP_SO_KEEP_ALIVE = "oak.mongo.socketKeepAlive";
+
     @Property(value = DEFAULT_URI,
             label = "Mongo URI",
             description = "Mongo connection URI used to connect to Mongo. Refer to " +
@@ -160,6 +166,14 @@ public class DocumentNodeStoreService {
     )
     private static final String PROP_DB = "db";
 
+    @Property(boolValue = DEFAULT_SO_KEEP_ALIVE,
+            label = "MongoDB socket keep-alive option",
+            description = "Whether socket keep-alive should be enabled for " +
+                    "connections to MongoDB. Note that this value can be " +
+                    "overridden via framework property 'oak.mongo.socketKeepAlive'"
+    )
+    static final String PROP_SO_KEEP_ALIVE = "socketKeepAlive";
+
     @Property(intValue = DEFAULT_CACHE,
             label = "Cache Size (in MB)",
             description = "Cache size in MB. This is distributed among various caches used in DocumentNodeStore"
@@ -440,6 +454,7 @@ private void registerNodeStoreIfPossible() throws IOException {
     private void registerNodeStore() throws IOException {
         String uri = PropertiesUtil.toString(prop(PROP_URI, FWK_PROP_URI), DEFAULT_URI);
         String db = PropertiesUtil.toString(prop(PROP_DB, FWK_PROP_DB), DEFAULT_DB);
+        boolean soKeepAlive = PropertiesUtil.toBoolean(prop(PROP_SO_KEEP_ALIVE, FWK_PROP_SO_KEEP_ALIVE), DEFAULT_SO_KEEP_ALIVE);
 
         int cacheSize = toInteger(prop(PROP_CACHE), DEFAULT_CACHE);
         int nodeCachePercentage = toInteger(prop(PROP_NODE_CACHE_PERCENTAGE), DEFAULT_NODE_CACHE_PERCENTAGE);
@@ -535,6 +550,7 @@ public void handleLeaseFailure() {
             }
 
             mkBuilder.setMaxReplicationLag(maxReplicationLagInSecs, TimeUnit.SECONDS);
+            mkBuilder.setSocketKeepAlive(soKeepAlive);
             mkBuilder.setMongoDB(uri, db, blobCacheSize);
 
             log.info("Connected to database '{}'", db);

oak-core/src/main/java/org/apache/jackrabbit/oak/plugins/document/util/MongoConnection.java

@@ -16,7 +16,6 @@
  */
 package org.apache.jackrabbit.oak.plugins.document.util;
 
-import java.net.UnknownHostException;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
@@ -28,6 +27,7 @@
 import com.mongodb.MongoClient;
 import com.mongodb.MongoClientOptions;
 import com.mongodb.MongoClientURI;
+import com.mongodb.MongoException;
 import com.mongodb.ReadConcern;
 import com.mongodb.ReadConcernLevel;
 import com.mongodb.WriteConcern;
@@ -50,10 +50,22 @@ public class MongoConnection {
      * See also http://docs.mongodb.org/manual/reference/connection-string/
      *
      * @param uri the MongoDB URI
-     * @throws UnknownHostException
+     * @throws MongoException if there are failures
      */
-    public MongoConnection(String uri) throws UnknownHostException  {
-        MongoClientOptions.Builder builder = MongoConnection.getDefaultBuilder();
+    public MongoConnection(String uri) throws MongoException {
+        this(uri, MongoConnection.getDefaultBuilder());
+    }
+
+    /**
+     * Constructs a new connection using the specified MongoDB connection
+     * String. The default client options are taken from the provided builder.
+     *
+     * @param uri the connection URI.
+     * @param builder the client option defaults.
+     * @throws MongoException if there are failures
+     */
+    public MongoConnection(String uri, MongoClientOptions.Builder builder)
+            throws MongoException {
         mongoURI = new MongoClientURI(uri, builder);
         mongo = new MongoClient(mongoURI);
     }
@@ -64,9 +76,10 @@ public MongoConnection(String uri) throws UnknownHostException  {
      * @param host The host address.
      * @param port The port.
      * @param database The database name.
-     * @throws Exception If an error occurred while trying to connect.
+     * @throws MongoException if there are failures
      */
-    public MongoConnection(String host, int port, String database) throws Exception {
+    public MongoConnection(String host, int port, String database)
+            throws MongoException {
         this("mongodb://" + host + ":" + port + "/" + database);
     }
 

oak-core/src/test/java/org/apache/jackrabbit/oak/plugins/document/DocumentNodeStoreServiceTest.java

@@ -17,11 +17,15 @@
 package org.apache.jackrabbit.oak.plugins.document;
 
 import java.io.File;
+import java.lang.reflect.Field;
 import java.util.Map;
 
 import com.google.common.collect.Maps;
+import com.mongodb.DB;
 
 import org.apache.commons.io.FilenameUtils;
+import org.apache.jackrabbit.oak.plugins.document.mongo.MongoDocumentStore;
+import org.apache.jackrabbit.oak.plugins.document.mongo.MongoDocumentStoreHelper;
 import org.apache.jackrabbit.oak.plugins.document.spi.JournalPropertyService;
 import org.apache.jackrabbit.oak.spi.state.NodeStore;
 import org.apache.jackrabbit.oak.stats.StatisticsProvider;
@@ -37,6 +41,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
 import static org.mockito.Mockito.mock;
 
@@ -134,6 +139,28 @@ public void setUpdateLimit() throws Exception {
         assertEquals(17, store.getUpdateLimit());
     }
 
+    @Test
+    public void keepAlive() throws Exception {
+        Map<String, Object> config = newConfig(repoHome);
+        config.put(DocumentNodeStoreService.PROP_SO_KEEP_ALIVE, true);
+        MockOsgi.activate(service, context.bundleContext(), config);
+        DocumentNodeStore store = context.getService(DocumentNodeStore.class);
+        MongoDocumentStore mds = getMongoDocumentStore(store);
+        DB db = MongoDocumentStoreHelper.getDB(mds);
+        assertTrue(db.getMongo().getMongoOptions().isSocketKeepAlive());
+    }
+
+    private static MongoDocumentStore getMongoDocumentStore(DocumentNodeStore s) {
+        try {
+            Field f = s.getClass().getDeclaredField("nonLeaseCheckingStore");
+            f.setAccessible(true);
+            return (MongoDocumentStore) f.get(s);
+        } catch (Exception e) {
+            fail(e.getMessage());
+            return null;
+        }
+    }
+
     private void assertPersistentCachePath(String expectedPath,
                                            String persistentCache,
                                            String repoHome) {

oak-core/src/test/java/org/apache/jackrabbit/oak/plugins/document/mongo/MongoConnectionTest.java

@@ -18,16 +18,19 @@
 
 import com.mongodb.DB;
 import com.mongodb.Mongo;
+import com.mongodb.MongoClientOptions;
 import com.mongodb.ReadConcern;
 import com.mongodb.ReplicaSetStatus;
 import com.mongodb.WriteConcern;
 
+import org.apache.jackrabbit.oak.plugins.document.MongoUtils;
 import org.apache.jackrabbit.oak.plugins.document.util.MongoConnection;
 import org.junit.Test;
 
 import static junitx.framework.ComparableAssert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -83,6 +86,26 @@ public void sufficientReadConcern() throws Exception {
         sufficientReadConcernSingleNode(ReadConcern.MAJORITY, true);
     }
 
+    @Test
+    public void socketKeepAlive() throws Exception {
+        assumeTrue(MongoUtils.isAvailable());
+        MongoClientOptions.Builder options = MongoConnection.getDefaultBuilder();
+        options.socketKeepAlive(true);
+        MongoConnection c = new MongoConnection(MongoUtils.URL, options);
+        try {
+            assertTrue(c.getDB().getMongo().getMongoOptions().isSocketKeepAlive());
+        } finally {
+            c.close();
+        }
+        // default is without keep-alive
+        c = new MongoConnection(MongoUtils.URL);
+        try {
+            assertFalse(c.getDB().getMongo().getMongoOptions().isSocketKeepAlive());
+        } finally {
+            c.close();
+        }
+    }
+
     private void sufficientWriteConcernReplicaSet(WriteConcern w,
                                                   boolean sufficient) {
         sufficientWriteConcern(w, true, sufficient);

oak-core/src/test/java/org/apache/jackrabbit/oak/plugins/document/mongo/MongoDocumentStoreHelper.java

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.jackrabbit.oak.plugins.document.mongo;
+
+import com.mongodb.DB;
+
+import org.apache.jackrabbit.oak.plugins.document.Collection;
+
+public final class MongoDocumentStoreHelper {
+
+    public static DB getDB(MongoDocumentStore store) {
+        return store.getDBCollection(Collection.NODES).getDB();
+    }
+}