From 53dc43bfe2fd9d3871844a1515f0e1c0fdf3c1a6 Mon Sep 17 00:00:00 2001 From: Steffen Pingel Date: Wed, 27 Nov 2024 09:53:12 +0100 Subject: [PATCH] Prefer mongodb.keep_alive_ms over spark.mongodb.keep_alive_ms property MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The documentation at https://www.mongodb.com/docs/spark-connector/master/faq/ specifies to set mongodb.keep_alive_ms but in the code the system property is prefixed with spark. This changes makes the code consistent with the documentation while maintaining backwards compatibility. When attempting to set this with --conf spark.executor.extraJavaOptions="-Dspark.mongodb.keep_alive_ms=70000" an exception occurs since Spark does not allow setting spark flags with that mechanism on executors: java.lang.Exception: spark.executor.extraJavaOptions is not allowed to set Spark options (was '-Dspark.mongodb.keep_alive_ms=70000’). Set them directly on a SparkConf or in a properties file when using ./bin/spark-submit. Work-around: Rather than setting the system property via extraJavaOptions on executors, the environment variable JAVA_TOOL_OPTIONS can be used: --conf spark.driver.extraJavaOptions=-Dspark.mongodb.keep_alive_ms=70000 --conf spark.executorEnv.JAVA_TOOL_OPTIONS=-Dspark.mongodb.keep_alive_ms=70000 --- .../connection/LazyMongoClientCache.java | 28 ++++++-- .../connection/LazyMongoClientCacheTest.java | 65 +++++++++++++++++++ 2 files changed, 87 insertions(+), 6 deletions(-) create mode 100644 src/test/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCacheTest.java diff --git a/src/main/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCache.java b/src/main/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCache.java index d96b49a8..fc2d0994 100644 --- a/src/main/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCache.java +++ b/src/main/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCache.java @@ -31,19 +31,35 @@ public final class LazyMongoClientCache { private static final MongoClientCache CLIENT_CACHE; - private static final String SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY = + static final String LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY = "spark.mongodb.keep_alive_ms"; + static final String SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY = "mongodb.keep_alive_ms"; + static { - int keepAliveMS = 5000; + int keepAliveMS = computeKeepAlive(5000); + + CLIENT_CACHE = new MongoClientCache(keepAliveMS); + } + + static int computeKeepAlive(int defaultValue) { + int keepAliveMS = defaultValue; try { - keepAliveMS = - Integer.parseInt(System.getProperty(SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, "5000")); + String legacyKeepAliveMS = + System.getProperty(LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY); + if (legacyKeepAliveMS != null) { + keepAliveMS = Integer.parseInt(legacyKeepAliveMS); + } } catch (NumberFormatException e) { // ignore and use default } - - CLIENT_CACHE = new MongoClientCache(keepAliveMS); + try { + keepAliveMS = Integer.parseInt(System.getProperty( + SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, Integer.toString(defaultValue))); + } catch (NumberFormatException e) { + // ignore and use default + } + return keepAliveMS; } /** diff --git a/src/test/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCacheTest.java b/src/test/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCacheTest.java new file mode 100644 index 00000000..a5340e98 --- /dev/null +++ b/src/test/java/com/mongodb/spark/sql/connector/connection/LazyMongoClientCacheTest.java @@ -0,0 +1,65 @@ +package com.mongodb.spark.sql.connector.connection; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class LazyMongoClientCacheTest { + + private String legacyPropertyValue; + private String propertyValue; + + @BeforeEach + void saveSystemProperties() { + legacyPropertyValue = + System.getProperty(LazyMongoClientCache.LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY); + propertyValue = + System.getProperty(LazyMongoClientCache.SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY); + } + + @AfterEach + void restoreSystemProperties() { + System.setProperty( + LazyMongoClientCache.LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, + (legacyPropertyValue != null) ? legacyPropertyValue : ""); + System.setProperty( + LazyMongoClientCache.SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, + (propertyValue != null) ? propertyValue : ""); + } + + @Test + void computeKeepAliveUsesDefaultIfPropertiesNotSet() { + System.setProperty(LazyMongoClientCache.LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, ""); + System.setProperty(LazyMongoClientCache.SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, ""); + + assertEquals(5000, LazyMongoClientCache.computeKeepAlive(5000)); + } + + @Test + void computeKeepAliveUsesLegacySystemProperty() { + System.setProperty( + LazyMongoClientCache.LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, "123"); + System.setProperty(LazyMongoClientCache.SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, ""); + + assertEquals(123, LazyMongoClientCache.computeKeepAlive(5000)); + } + + @Test + void computeKeepAliveUsesMongoSystemProperty() { + System.setProperty(LazyMongoClientCache.LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, ""); + System.setProperty(LazyMongoClientCache.SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, "123"); + + assertEquals(123, LazyMongoClientCache.computeKeepAlive(5000)); + } + + @Test + void computeKeepAlivePrefersMongoSystemProperty() { + System.setProperty( + LazyMongoClientCache.LEGACY_SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, "456"); + System.setProperty(LazyMongoClientCache.SYSTEM_MONGO_CACHE_KEEP_ALIVE_MS_PROPERTY, "123"); + + assertEquals(123, LazyMongoClientCache.computeKeepAlive(5000)); + } +}