fixes #185: Consumer not working in Cluster environment

conker84 · mneedham · commit 6bd66a62f376 · 2019-06-05T15:38:20.000+01:00
diff --git a/common/src/main/kotlin/streams/utils/Neo4jUtils.kt b/common/src/main/kotlin/streams/utils/Neo4jUtils.kt
@@ -1,11 +1,19 @@
 package streams.utils
 
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
 import org.neo4j.graphdb.QueryExecutionException
 import org.neo4j.kernel.internal.GraphDatabaseAPI
 import org.neo4j.logging.internal.LogService
+import org.neo4j.logging.Log
+import org.neo4j.logging.NullLog
 import java.lang.reflect.InvocationTargetException
+import kotlin.streams.toList
 
 object Neo4jUtils {
+    @JvmStatic val LEADER = "LEADER"
     fun isWriteableInstance(db: GraphDatabaseAPI): Boolean {
         try {
             val isSlave = StreamsUtils.ignoreExceptions(
@@ -19,7 +27,7 @@ object Neo4jUtils {
             }
 
             val role = db.execute("CALL dbms.cluster.role()").columnAs<String>("role").next()
-            return role.equals("LEADER", ignoreCase = true)
+            return role.equals(LEADER, ignoreCase = true)
         } catch (e: QueryExecutionException) {
             if (e.statusCode.equals("Neo.ClientError.Procedure.ProcedureNotFound", ignoreCase = true)) {
                 return true
@@ -32,4 +40,55 @@ object Neo4jUtils {
         return db.dependencyResolver
                 .resolveDependency(LogService::class.java)
     }
+
+    fun isEnterpriseEdition(db: GraphDatabaseAPI): Boolean {
+        try {
+            return db.execute("""
+                CALL dbms.components() YIELD edition
+                RETURN edition = "enterprise" AS isEnterprise
+            """.trimIndent()).columnAs<Boolean>("isEnterprise").next()
+        } catch (e: QueryExecutionException) {
+            if (e.statusCode.equals("Neo.ClientError.Procedure.ProcedureNotFound", ignoreCase = true)) {
+                return false
+            }
+            throw e
+        }
+    }
+
+    fun clusterHasLeader(db: GraphDatabaseAPI): Boolean {
+        try {
+            return db.execute("""
+                CALL dbms.cluster.overview() YIELD role
+                RETURN role
+            """.trimIndent())
+                    .columnAs<String>("role")
+                    .stream()
+                    .toList()
+                    .contains(LEADER)
+        } catch (e: QueryExecutionException) {
+            if (e.statusCode.equals("Neo.ClientError.Procedure.ProcedureNotFound", ignoreCase = true)) {
+                return false
+            }
+            throw e
+        }
+    }
+
+    fun <T> executeInWriteableInstance(db: GraphDatabaseAPI, action: () -> T) {
+        if (isWriteableInstance(db)) {
+            action()
+        }
+    }
+
+    fun executeInLeader(db: GraphDatabaseAPI, log: Log, timeout: Long = 120000, action: () -> Unit) {
+        GlobalScope.launch(Dispatchers.IO) {
+            val start = System.currentTimeMillis()
+            val delay: Long = 2000
+            while (!clusterHasLeader(db) && System.currentTimeMillis() - start < timeout) {
+                log.info("$LEADER not found, new check comes in $delay milliseconds...")
+                delay(delay)
+            }
+            executeInWriteableInstance(db, action)
+        }
+    }
+
 }
diff --git a/common/src/test/kotlin/streams/utils/Neo4jUtilsTest.kt b/common/src/test/kotlin/streams/utils/Neo4jUtilsTest.kt
@@ -5,6 +5,7 @@ import org.junit.BeforeClass
 import org.junit.Test
 import org.neo4j.kernel.internal.GraphDatabaseAPI
 import org.neo4j.test.TestGraphDatabaseFactory
+import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 
 class Neo4jUtilsTest {
@@ -32,4 +33,10 @@ class Neo4jUtilsTest {
         assertTrue { isWriteableInstance }
     }
 
+    @Test
+    fun shouldCheckIfIsEnterpriseEdition() {
+        val isEnterprise = Neo4jUtils.isEnterpriseEdition(db)
+        assertFalse { isEnterprise }
+    }
+
 }
diff --git a/consumer/src/main/kotlin/streams/StreamsEventSink.kt b/consumer/src/main/kotlin/streams/StreamsEventSink.kt
@@ -1,12 +1,14 @@
 package streams
 
 import org.neo4j.kernel.configuration.Config
+import org.neo4j.kernel.internal.GraphDatabaseAPI
 import org.neo4j.logging.Log
 
 abstract class StreamsEventSink(private val config: Config,
                                 private val queryExecution: StreamsEventSinkQueryExecution,
                                 private val streamsTopicService: StreamsTopicService,
-                                private val log: Log) {
+                                private val log: Log,
+                                private val db: GraphDatabaseAPI) {
 
     abstract fun stop()
 
@@ -36,13 +38,14 @@ abstract class StreamsEventConsumerFactory {
 
 object StreamsEventSinkFactory {
     fun getStreamsEventSink(config: Config, streamsQueryExecution: StreamsEventSinkQueryExecution,
-                            streamsTopicService: StreamsTopicService, log: Log): StreamsEventSink {
+                            streamsTopicService: StreamsTopicService, log: Log, db: GraphDatabaseAPI): StreamsEventSink {
         return Class.forName(config.raw.getOrDefault("streams.sink", "streams.kafka.KafkaEventSink"))
                 .getConstructor(Config::class.java,
                         StreamsEventSinkQueryExecution::class.java,
                         StreamsTopicService::class.java,
-                        Log::class.java)
-                .newInstance(config, streamsQueryExecution, streamsTopicService, log) as StreamsEventSink
+                        Log::class.java,
+                        GraphDatabaseAPI::class.java)
+                .newInstance(config, streamsQueryExecution, streamsTopicService, log, db) as StreamsEventSink
     }
 }
 
diff --git a/consumer/src/main/kotlin/streams/StreamsEventSinkExtensionFactory.kt b/consumer/src/main/kotlin/streams/StreamsEventSinkExtensionFactory.kt
@@ -53,28 +53,30 @@ class StreamsEventSinkExtensionFactory : KernelExtensionFactory<StreamsEventSink
                                 log.getUserLog(StreamsEventSinkQueryExecution::class.java),
                                 strategyMap)
 
-                        // Create and start the Sink
+                        // Create the Sink
+                        val sinkLog = log.getUserLog(StreamsEventSinkFactory::class.java)
                         eventSink = StreamsEventSinkFactory
                                 .getStreamsEventSink(configuration,
                                         streamsQueryExecution,
                                         streamsTopicService,
-                                        log.getUserLog(StreamsEventSinkFactory::class.java))
-                        eventSink.start()
-                        if (Neo4jUtils.isWriteableInstance(db)) {
-                            if (streamsLog.isDebugEnabled) {
-                                streamsLog.debug("Subscribed topics with Cypher queries: ${streamsTopicService.getAllCypherTemplates()}")
-                                streamsLog.debug("Subscribed topics with CDC configuration: ${streamsTopicService.getAllCDCTopics()}")
-                            } else {
-                                streamsLog.info("Subscribed topics: ${streamsTopicService.getTopics()}")
-                            }
+
+                                        sinkLog,
+                                        db)
+                        // start the Sink
+                        if (Neo4jUtils.isEnterpriseEdition(db)) {
+                            sinkLog.info("The Sink module is running in an enterprise edition, checking for the ${Neo4jUtils.LEADER}")
+                            Neo4jUtils.executeInLeader(db, sinkLog) { initSinkModule() }
+                        } else {
+                            // check if is writeable instance
+                            Neo4jUtils.executeInWriteableInstance(db) { initSinkModule() }
                         }
 
                         // Register required services for the Procedures
                         StreamsSinkProcedures.registerStreamsSinkConfiguration(streamsSinkConfiguration)
                         eventSink = StreamsEventSinkFactory.getStreamsEventSink(configuration,
                                 streamsQueryExecution,
                                 streamsTopicService,
-                                log.getUserLog(StreamsEventSinkFactory::class.java))
+                                log.getUserLog(StreamsEventSinkFactory::class.java), db)
                         eventSink.start()
                         StreamsSinkProcedures.registerStreamsEventConsumerFactory(eventSink.getEventConsumerFactory())
                         StreamsSinkProcedures.registerStreamsEventSinkConfigMapper(eventSink.getEventSinkConfigMapper())
@@ -87,6 +89,11 @@ class StreamsEventSinkExtensionFactory : KernelExtensionFactory<StreamsEventSink
             }
         }
 
+        private fun initSinkModule() {
+            eventSink.start()
+            streamsLog.info("Streams Sink module initialised")
+        }
+
         override fun stop() {
             try {
                 StreamsUtils.ignoreExceptions({ eventSink.stop() }, UninitializedPropertyAccessException::class.java)
diff --git a/consumer/src/main/kotlin/streams/StreamsSinkConfiguration.kt b/consumer/src/main/kotlin/streams/StreamsSinkConfiguration.kt
@@ -6,7 +6,7 @@ import streams.service.sink.strategy.SourceIdIngestionStrategyConfig
 import streams.service.Topics
 
 
-private object StreamsSinkConfigurationConstants {
+object StreamsSinkConfigurationConstants {
     const val STREAMS_CONFIG_PREFIX: String = "streams."
     const val ENABLED = "sink.enabled"
     const val PROCEDURES_ENABLED = "procedures.enabled"
diff --git a/consumer/src/main/kotlin/streams/StreamsTopicService.kt b/consumer/src/main/kotlin/streams/StreamsTopicService.kt
@@ -18,7 +18,7 @@ class StreamsTopicService(private val db: GraphDatabaseAPI) {
     private val properties: GraphProperties = db.dependencyResolver.resolveDependency(EmbeddedProxySPI::class.java).newGraphPropertiesProxy()
     private val log = Neo4jUtils.getLogService(db).getUserLog(StreamsTopicService::class.java)
 
-    fun clearAll() {
+    fun clearAll() { // TODO move to Neo4jUtils#executeInWriteableInstance
         if (!Neo4jUtils.isWriteableInstance(db)) {
             return
         }
diff --git a/consumer/src/main/kotlin/streams/kafka/KafkaEventSink.kt b/consumer/src/main/kotlin/streams/kafka/KafkaEventSink.kt
@@ -4,18 +4,21 @@ import kotlinx.coroutines.*
 import org.apache.kafka.clients.consumer.*
 import org.apache.kafka.common.TopicPartition
 import org.neo4j.kernel.configuration.Config
+import org.neo4j.kernel.internal.GraphDatabaseAPI
 import org.neo4j.logging.Log
 import streams.*
-import streams.kafka.KafkaTopicConfig.Companion.toTopicPartitionMap
 import streams.serialization.JSONUtils
+import streams.utils.Neo4jUtils
 import streams.utils.StreamsUtils
 import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.TimeUnit
 
 
 class KafkaEventSink(private val config: Config,
                      private val queryExecution: StreamsEventSinkQueryExecution,
                      private val streamsTopicService: StreamsTopicService,
-                     private val log: Log): StreamsEventSink(config, queryExecution, streamsTopicService, log) {
+                     private val log: Log,
+                     private val db: GraphDatabaseAPI): StreamsEventSink(config, queryExecution, streamsTopicService, log, db) {
 
     private lateinit var eventConsumer: StreamsEventConsumer
     private lateinit var job: Job
@@ -46,15 +49,28 @@ class KafkaEventSink(private val config: Config,
 
     override fun start() {
         val streamsConfig = StreamsSinkConfiguration.from(config)
+        val topics = streamsTopicService.getTopics()
+        val isWriteableInstance = Neo4jUtils.isWriteableInstance(db)
         if (!streamsConfig.enabled) {
+            if (topics.isNotEmpty() && isWriteableInstance) {
+                log.warn("You configured the following topics: $topics, in order to make the Sink work please set the `${StreamsSinkConfigurationConstants.STREAMS_CONFIG_PREFIX}${StreamsSinkConfigurationConstants.ENABLED}` to `true`")
+            }
             return
         }
         log.info("Starting the Kafka Sink")
         this.eventConsumer = getEventConsumerFactory()
                 .createStreamsEventConsumer(config.raw, log)
-                .withTopics(streamsTopicService.getTopics())
+                .withTopics(topics)
         this.eventConsumer.start()
         this.job = createJob()
+        if (isWriteableInstance) {
+            if (log.isDebugEnabled) {
+                log.debug("Subscribed topics with Cypher queries: ${streamsTopicService.getAllCypherTemplates()}")
+                log.debug("Subscribed topics with CDC configuration: ${streamsTopicService.getAllCDCTopics()}")
+            } else {
+                log.info("Subscribed topics: $topics")
+            }
+        }
         log.info("Kafka Sink started")
     }
 
@@ -82,11 +98,19 @@ class KafkaEventSink(private val config: Config,
         return GlobalScope.launch(Dispatchers.IO) { // TODO improve exception management
             try {
                 while (isActive) {
-                    eventConsumer.read { topic, data ->
+                    if (Neo4jUtils.isWriteableInstance(db)) {
+                        eventConsumer.read { topic, data ->
+                            if (log.isDebugEnabled) {
+                                log.debug("Reading data from topic $topic")
+                            }
+                            queryExecution.writeForTopic(topic, data)
+                        }
+                    } else {
+                        val timeMillis = TimeUnit.MILLISECONDS.toMinutes(5)
                         if (log.isDebugEnabled) {
-                            log.debug("Reading data from topic $topic")
+                            log.debug("Not in a writeable instance, new check in $timeMillis millis")
                         }
-                        queryExecution.writeForTopic(topic, data)
+                        delay(timeMillis)
                     }
                 }
                 eventConsumer.stop()

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ class StreamsTopicService(private val db: GraphDatabaseAPI) {`
`18`	`18`	`private val properties: GraphProperties = db.dependencyResolver.resolveDependency(EmbeddedProxySPI::class.java).newGraphPropertiesProxy()`
`19`	`19`	`private val log = Neo4jUtils.getLogService(db).getUserLog(StreamsTopicService::class.java)`
`20`	`20`
`21`		`- fun clearAll() {`
	`21`	`+ fun clearAll() { // TODO move to Neo4jUtils#executeInWriteableInstance`
`22`	`22`	`if (!Neo4jUtils.isWriteableInstance(db)) {`
`23`	`23`	`return`
`24`	`24`	`}`