Merge pull request boskworks#235 from prdoyle/timeouts

Revamp Mongo timeouts

Author: prdoyle

bosk-mongo/src/main/java/works/bosk/drivers/mongo/internal/ChangeReceiver.java

@@ -15,6 +15,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicLong;
 import org.bson.BsonDocument;
+import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.slf4j.MDC;
@@ -56,6 +57,7 @@ class ChangeReceiver implements Closeable {
 	private final MongoCollection<BsonDocument> collection;
 	private final ScheduledExecutorService ex = Executors.newScheduledThreadPool(1);
 	private final Exception creationPoint;
+	private volatile @Nullable Thread thread = null;
 	private volatile boolean isClosed = false;
 
 	ChangeReceiver(String boskName, Identifier boskID, ChangeListener listener, MongoDriverSettings settings, MongoCollection<BsonDocument> collection) {
@@ -70,6 +72,7 @@ class ChangeReceiver implements Closeable {
 			// to ensure the database is set up for change streams.
 			probeChangeStreamCursor();
 		}
+		LOGGER.debug("Scheduling ChangeReceiver connectionLoop task with {} ms interval", settings.timescaleMS());
 		ex.scheduleWithFixedDelay(
 			this::connectionLoop,
 			0,
@@ -78,6 +81,23 @@ class ChangeReceiver implements Closeable {
 		);
 	}
 
+	/**
+	 * If the connectionLoop is running, interrupt it.
+	 * (Otherwise it's going to reinitialize on its own anyway,
+	 * so there's no need to do anything because the effect is the same.)
+	 */
+	public void interrupt() {
+		Thread t = thread;
+		if (t == null) {
+			LOGGER.debug("ChangeReceiver thread is not running; no need to interrupt");
+		} else if (t == currentThread()) {
+			LOGGER.debug("ChangeReceiver thread is doing the disconnecting; no need to interrupt");
+		} else {
+			LOGGER.debug("Interrupting ChangeReceiver thread {}", t.getName());
+			t.interrupt();
+		}
+	}
+
 	private void probeChangeStreamCursor() {
 		try (var _ = openCursor()) {
 			LOGGER.debug("Successfully opened MongoDB cursor");
@@ -109,6 +129,7 @@ private void connectionLoop() {
 		try (MDCScope _ = setupMDC(boskName, boskID)) {
 			LOGGER.debug("Starting connectionLoop task");
 			try {
+				thread = currentThread();
 				while (!isClosed) {
 					// Design notes:
 					//
@@ -197,6 +218,7 @@ private void connectionLoop() {
 			} finally {
 				LOGGER.debug("Ending connectionLoop task; isClosed={}", isClosed);
 				currentThread().setName(oldThreadName);
+				thread = null;
 			}
 		} catch (RuntimeException e) {
 			addContextToException(e);
@@ -221,7 +243,6 @@ private void addContextToException(Throwable x) {
 	private MongoChangeStreamCursor<ChangeStreamDocument<BsonDocument>> openCursor() {
 		MongoChangeStreamCursor<ChangeStreamDocument<BsonDocument>> result = collection
 			.watch()
-			.maxAwaitTime(settings.timescaleMS(), MILLISECONDS)
 			.cursor();
 		LOGGER.debug("Cursor is open");
 		return result;

bosk-mongo/src/main/java/works/bosk/drivers/mongo/internal/MainDriver.java

@@ -78,8 +78,8 @@ public final class MainDriver<R extends StateTreeNode> implements MongoDriver {
 	final Formatter formatter;
 
 	final long flushTimeout;
-	final long initialRootTimeout;
-	final long driverReinitializeTimeout;
+	final long initializeTimeout;
+	final long reinitializationTimeout;
 
 	/**
 	 * {@link MongoClient#close()} throws if called more than once.
@@ -123,26 +123,77 @@ public MainDriver(
 			this.bsonSerializer = bsonSerializer;
 			this.downstream = downstream;
 
-			MongoClientSettings.Builder builder = MongoClientSettings
-				.builder(clientSettings);
+			// Flushes work by waiting for the latest version to arrive on the change stream.
+			// If we wait for two heartbeats and don't see the update, something has gone wrong.
+			//
+			// (Note that flush does a retry, so it will actually take
+			// twice as long as this before throwing a FlushFailureException.)
+			flushTimeout = 2L * driverSettings.timescaleMS();
+
+			// Initialization must wait for a heartbeat to succeed, so we wait twice that long,
+			// plus one more for the network connection and initial query.
+			initializeTimeout = 3L * driverSettings.timescaleMS();
+
+			// The sum of the steps required to reinitialize after a disconnect,
+			// plus a little extra to make sure we don't cut it off when it's about to succeed.
+			reinitializationTimeout =
+				2L * driverSettings.timescaleMS() // ChangeStream reconnect
+					+ initializeTimeout // Initialize after reconnecting
+					+ driverSettings.timescaleMS() // Extra buffer
+			;
+
+			MongoClientSettings.Builder commonSettingsBuilder = MongoClientSettings
+				.builder(clientSettings)
+				.applyToServerSettings(s ->
+					// If timescaleMS is shorter than the default min heartbeat,
+					// then we need to reduce this setting to prevent the client
+					// from using a stale view of the server state for too long.
+					// If timescaleMS is longer, then the user has told us
+					// they don't mind longer delays and want the increased
+					// efficiency of fewer heartbeats.
+					// Either way, timescaleMS is the right value for this setting.
+					//
+					// Note that this doesn't set the heartbeat frequency itself.
+					// That is left at the default value, since it is only used
+					// to "notice" connectivity problems when the driver is quiescent,
+					// which is not time-critical and is not governed by timescaleMS:
+					// the actual behaviour of the bosk during a network partition
+					// is that its contents remain fixed, and it doesn't matter much
+					// whether that is achieved by formally disconnecting or simply
+					// by doing nothing.
+					s.minHeartbeatFrequency(driverSettings.timescaleMS(), MILLISECONDS))
+				;
 
 			// By default, we deal only with durable data that won't get rolled back.
 			// In some circumstances, we need the very latest possible data for correctness,
 			// so we override the ReadConcern in those cases.
-			builder
+			commonSettingsBuilder
 				.readConcern(ReadConcern.MAJORITY)
 				.writeConcern(WriteConcern.MAJORITY);
 
+			var changeStreamSettingsBuilder = MongoClientSettings.builder(commonSettingsBuilder.build())
+				.applyToClusterSettings(c ->
+					c.serverSelectionTimeout(initializeTimeout, MILLISECONDS))
+				.applyToSocketSettings(s ->
+					s.connectTimeout(initializeTimeout, MILLISECONDS)
+						// No read timeout for change streams; they can be idle indefinitely
+						.readTimeout(0, MILLISECONDS))
+				;
+
+			var changeStreamClient = MongoClients.create(changeStreamSettingsBuilder.build());
+			closeables.addFirst(changeStreamClient);
+
 			// Override timeouts to make them compatible with driverSettings.timescaleMS()
-			builder
-				.timeout(2L * driverSettings.timescaleMS(), MILLISECONDS);
+			var querySettingsBuilder = MongoClientSettings.builder(commonSettingsBuilder.build());
+			querySettingsBuilder
+				.timeout(flushTimeout, MILLISECONDS);
+
+			var queryClient = MongoClients.create(querySettingsBuilder.build());
+			closeables.addFirst(queryClient);
 
-			var mongoClient = MongoClients.create(builder.build());
-			closeables.addFirst(mongoClient);
-			MongoCollection<BsonDocument> changeStreamCollection = mongoClient
+			this.queryCollection = TransactionalCollection.of(queryClient
 				.getDatabase(driverSettings.database())
-				.getCollection(COLLECTION_NAME, BsonDocument.class);
-			this.queryCollection = TransactionalCollection.of(changeStreamCollection, mongoClient);
+				.getCollection(COLLECTION_NAME, BsonDocument.class), queryClient);
 			LOGGER.debug("Using database \"{}\" collection \"{}\"", driverSettings.database(), COLLECTION_NAME);
 
 			this.formatter = new Formatter(boskInfo, bsonSerializer);
@@ -153,20 +204,13 @@ public MainDriver(
 			if (factory != null) {
 				listener = factory.apply(listener);
 			}
+
+			MongoCollection<BsonDocument> changeStreamCollection = changeStreamClient
+				.getDatabase(driverSettings.database())
+				.getCollection(COLLECTION_NAME, BsonDocument.class);
 			this.receiver = new ChangeReceiver(boskInfo.name(), boskInfo.instanceID(), listener, driverSettings, changeStreamCollection);
 		}
 
-		// Flushes work by waiting for the latest version to arrive on the change stream.
-		// If we wait twice as long as that takes, and we don't see the update, something
-		// has gone wrong.
-		flushTimeout = 2L * driverSettings.timescaleMS();
-
-		// TODO: Justify this calculation.
-		initialRootTimeout = 5L * driverSettings.timescaleMS();
-
-		// The ChangeStream resets itself after timescaleMS, so it needs
-		// several times that long to restore itself and publish a new driver.
-		driverReinitializeTimeout = 5L * driverSettings.timescaleMS();
 	}
 
 	@Override
@@ -478,7 +522,7 @@ public void onConnectionSucceeded() throws
 		private void runInitialRootAction(FutureTask<R> initialRootAction) throws InterruptedException, TimeoutException, InitialRootActionException {
 			initialRootAction.run();
 			try {
-				initialRootAction.get(initialRootTimeout, MILLISECONDS);
+				initialRootAction.get(initializeTimeout, MILLISECONDS);
 				LOGGER.debug("initialRoot action completed successfully");
 			} catch (ExecutionException e) {
 				LOGGER.debug("initialRoot action failed", e);
@@ -544,6 +588,7 @@ private FormatDriver<R> newPreferredFormatDriver() {
 	}
 
 	private FormatDriver<R> detectFormat() throws UninitializedCollectionException, UnrecognizedFormatException {
+		LOGGER.debug("Detecting format");
 		Manifest manifest = loadManifest();
 		DatabaseFormat format = manifest.pando().isPresent()? manifest.pando().get() : SEQUOIA;
 		BsonString documentId = (format == SEQUOIA)
@@ -621,7 +666,7 @@ private MDCScope beginDriverOperation(String description, Object... args) {
 			throw new IllegalStateException("Driver is closed");
 		}
 		MDCScope ex = setupMDC(boskInfo.name(), boskInfo.instanceID());
-		LOGGER.debug(description, args);
+		LOGGER.debug(description + " w/" + this.formatDriver.getClass().getSimpleName(), args);
 		if (driverSettings.testing().eventDelayMS() < 0) {
 			LOGGER.debug("| eventDelayMS {}ms ", driverSettings.testing().eventDelayMS());
 			try {
@@ -655,8 +700,9 @@ private <X extends Exception, Y extends Exception> void doRetryableDriverOperati
 							throw new DisconnectedException(e);
 						}
 					} else {
-						LOGGER.debug("MongoException is not recoverable; rethrowing", e);
-						throw e;
+						LOGGER.debug("MongoException is not recoverable; disconnecting", e);
+						setDisconnectedDriver(e);
+						throw new DisconnectedException(e);
 					}
 				}
 				break;
@@ -680,8 +726,8 @@ private <X extends Exception, Y extends Exception> void doRetryableDriverOperati
 	private <X extends Exception, Y extends Exception> void waitAndRetry(RetryableOperation<X, Y> operation, String description, Object... args) throws X, Y {
 		try {
 			formatDriverLock.lock();
-			LOGGER.debug("Waiting for new FormatDriver for {} ms", driverReinitializeTimeout);
-			boolean success = formatDriverChanged.await(driverReinitializeTimeout, MILLISECONDS);
+			LOGGER.debug("Waiting for new FormatDriver for {} ms", reinitializationTimeout);
+			boolean success = formatDriverChanged.await(reinitializationTimeout, MILLISECONDS);
 			if (!success) {
 				LOGGER.warn("Timed out waiting for MongoDB to recover; will retry anyway, but the operation may fail");
 			}
@@ -705,7 +751,7 @@ private <X extends Exception, Y extends Exception> void waitAndRetry(RetryableOp
 		} finally {
 			formatDriverLock.unlock();
 		}
-		LOGGER.debug("Retrying " + description, args);
+		LOGGER.debug("Retrying " + description + " w/" + this.formatDriver.getClass().getSimpleName(), args);
 		operation.run();
 	}
 
@@ -716,14 +762,22 @@ private <X extends Exception, Y extends Exception> void waitAndRetry(RetryableOp
 	 * better driver to arrive instead.
 	 */
 	void setDisconnectedDriver(Throwable reason) {
-		LOGGER.debug("quietlySetDisconnectedDriver({}) (previously {})", reason.getClass().getSimpleName(), formatDriver.getClass().getSimpleName());
+		LOGGER.debug("setDisconnectedDriver({}) (previously {})", reason.getClass().getSimpleName(), formatDriver.getClass().getSimpleName());
+		FormatDriver<R> oldDriver;
 		try {
 			formatDriverLock.lock();
-			formatDriver.close();
+			oldDriver = formatDriver;
+			oldDriver.close();
 			formatDriver = new DisconnectedDriver<>(reason);
 		} finally {
 			formatDriverLock.unlock();
 		}
+
+		if (!(oldDriver instanceof DisconnectedDriver<?>)) {
+			// The receiver is what reconnects us. Poke it to make sure it knows things
+			// have gone south, and we need to try to reconnect.
+			receiver.interrupt();
+		}
 	}
 
 	/**

bosk-mongo/src/test/java/works/bosk/drivers/mongo/internal/ErrorRecordingChangeListener.java

@@ -0,0 +1,46 @@
+package works.bosk.drivers.mongo.internal;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeoutException;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class ErrorRecordingChangeListener extends ForwardingChangeListener {
+	final ErrorRecorder errorRecorder;
+
+	public static class ErrorRecorder {
+		public int failureCount = 0;
+		public final List<Throwable> disconnections = new ArrayList<>();
+
+		public void assertAllClear(String description) {
+			assertTrue(failureCount == 0 && disconnections.isEmpty(),
+				"Expected no errors " + description + "; found " + this);
+		}
+
+		@Override
+		public String toString() {
+			return "ErrorRecorder{" +
+				"failureCount=" + failureCount +
+				", disconnections=" + disconnections +
+				'}';
+		}
+	}
+
+	public ErrorRecordingChangeListener(ErrorRecorder errorRecorder, ChangeListener downstream) {
+		super(downstream);
+		this.errorRecorder = errorRecorder;
+	}
+
+	@Override
+	public void onConnectionFailed() throws InterruptedException, TimeoutException {
+		errorRecorder.failureCount++;
+		super.onConnectionFailed();
+	}
+
+	@Override
+	public void onDisconnect(Throwable e) {
+		errorRecorder.disconnections.add(e);
+		super.onDisconnect(e);
+	}
+}

bosk-mongo/src/test/java/works/bosk/drivers/mongo/internal/ForwardingChangeListener.java

@@ -0,0 +1,34 @@
+package works.bosk.drivers.mongo.internal;
+
+import com.mongodb.client.model.changestream.ChangeStreamDocument;
+import java.io.IOException;
+import java.util.concurrent.TimeoutException;
+import org.bson.BsonDocument;
+
+class ForwardingChangeListener implements ChangeListener {
+	final ChangeListener downstream;
+
+	public ForwardingChangeListener(ChangeListener downstream) {
+		this.downstream = downstream;
+	}
+
+	@Override
+	public void onConnectionSucceeded() throws UnrecognizedFormatException, UninitializedCollectionException, InterruptedException, IOException, InitialRootActionException, TimeoutException, FailedSessionException {
+		downstream.onConnectionSucceeded();
+	}
+
+	@Override
+	public void onEvent(ChangeStreamDocument<BsonDocument> event) throws UnprocessableEventException {
+		downstream.onEvent(event);
+	}
+
+	@Override
+	public void onConnectionFailed() throws InterruptedException, TimeoutException {
+		downstream.onConnectionFailed();
+	}
+
+	@Override
+	public void onDisconnect(Throwable e) {
+		downstream.onDisconnect(e);
+	}
+}

bosk-mongo/src/test/java/works/bosk/drivers/mongo/internal/MongoDriverConformanceTest.java

@@ -34,6 +34,19 @@ class MongoDriverConformanceTest extends SharedDriverConformanceTest {
 	private static MongoService mongoService;
 	private final MongoDriverSettings driverSettings;
 	private final AtomicInteger numOpenDrivers = new AtomicInteger(0);
+	private ErrorRecordingChangeListener.ErrorRecorder errorRecorder;
+
+	@BeforeEach
+	void setupErrorRecording() {
+		errorRecorder = new ErrorRecordingChangeListener.ErrorRecorder();
+		MainDriver.LISTENER_FACTORY.set(downstream -> new ErrorRecordingChangeListener(errorRecorder, downstream));
+	}
+
+	@AfterEach
+	void teardownErrorRecording() {
+		errorRecorder.assertAllClear("after test");
+		MainDriver.LISTENER_FACTORY.remove();
+	}
 
 	public MongoDriverConformanceTest(ParameterSet parameters) {
 		this.driverSettings = parameters.driverSettingsBuilder().build();