Kafka issue with ClassLoader not containing the CodecWrapper class

src/com/workflowfm/pew/stateless/instances/kafka/components/Tracked.scala

@@ -8,6 +8,7 @@ import akka.stream.scaladsl.{Sink, Source}
 import akka.{Done, NotUsed}
 import com.workflowfm.pew.stateless.StatelessMessages.AnyMsg
 import com.workflowfm.pew.stateless.instances.kafka.settings.KafkaExecutorSettings
+import org.apache.kafka.clients.producer.KafkaProducer
 import org.bson.types.ObjectId
 
 import scala.concurrent.{ExecutionContext, Future}
@@ -136,6 +137,33 @@ object Tracked {
 
     consuming.tail.foldLeft[T[Seq[Msg]]]( foldStart )( combine )
   }
+
+  /** Execute a function on the current thread whilst using a specific ClassLoader.
+    * Restores the previous ClassLoader before returning.
+    *
+    * @param tmpClassLoader ClassLoader to use during the function execution.
+    * @param fnWrapped Function to execute with a different ClassLoader.
+    * @tparam T The return type of `fnWrapped`
+    * @return The returned value of `fnWrapped`
+    */
+  def withClassLoader[T]( tmpClassLoader: ClassLoader )( fnWrapped: => T ): T = {
+    val pushedClassLoader = Thread.currentThread().getContextClassLoader
+    try{
+      Thread.currentThread().setContextClassLoader( tmpClassLoader )
+      fnWrapped
+    } finally {
+      Thread.currentThread().setContextClassLoader( pushedClassLoader )
+    }
+  }
+
+  /** Create a KafkaProducer from settings whilst explicitly un-setting the ClassLoader.
+    * This by-passes errors encountered in the KafkaProducer constructor where Threads
+    * within an ExecutionContext do not list the necessary key or value serialiser classes.
+    * Explicitly setting `null` causes the constructor to use the Kafka ClassLoader
+    * which should contain these values.
+    */
+  def createProducer[K, V]( settings: ProducerSettings[K, V] ): KafkaProducer[K, V]
+    = withClassLoader( null ) { settings.createKafkaProducer() }
 }
 
 /** A Tracked type which uses a `Committable` as tracking information.
@@ -155,7 +183,7 @@ trait HasCommittableSinks[T[X] <: HasCommittable[X]]
   with TrackedMultiSink[T] {
 
   override def sink[V <: AnyMsg]( implicit s: KafkaExecutorSettings ): Sink[HasCommittable[V], Future[Done]]
-    = Producer.commitableSink( s.psAllMessages )
+    = Producer.commitableSink( s.psAllMessages, Tracked.createProducer( s.psAllMessages ) )
       .contramap( msg =>
         ProducerMessage.Message(
           s.record( msg.value ),
@@ -164,7 +192,7 @@ trait HasCommittableSinks[T[X] <: HasCommittable[X]]
       )
 
   override def sinkMulti[V <: AnyMsg]( implicit s: KafkaExecutorSettings ): Sink[HasCommittable[Seq[V]], Future[Done]]
-    = Producer.commitableSink( s.psAllMessages )
+    = Producer.commitableSink( s.psAllMessages, Tracked.createProducer( s.psAllMessages ) )
       .contramap( msgs =>
         ProducerMessage.MultiMessage(
           msgs.value.map(s.record).to,
@@ -215,7 +243,6 @@ object CommitTracked
           msg.committableOffset
         )
       )
-
 }
 
 /** A Tracked wrapper superclass for Tracked types which know the partition of their inputs.
@@ -330,6 +357,7 @@ object Transaction
       .map( msg => Transaction( msg.record.value(), msg.partitionOffset ) )
 
   override def sink[V <: AnyMsg]( implicit s: KafkaExecutorSettings ): Sink[Transaction[V], Future[Done]] = {
+    // TODO: Construct the KafkaProducer with the correct ClassLoader
     val id = ObjectId.get.toString
     Transactional.sink( s.psAllMessages, id )
     .contramap( msg =>
@@ -341,6 +369,7 @@ object Transaction
   }
 
   override def sinkMulti[V <: AnyMsg]( implicit s: KafkaExecutorSettings ): Sink[Transaction[Seq[V]], Future[Done]] = {
+    // TODO: Construct the KafkaProducer with the correct ClassLoader
     val id = ObjectId.get.toString
     Transactional.sink( s.psAllMessages, id )
     .contramap( msgs =>
@@ -385,7 +414,7 @@ object Untracked
       )
 
   override def sink[Value <: AnyMsg]( implicit s: KafkaExecutorSettings ): Sink[Untracked[Value], Future[Done]]
-    = Producer.plainSink( s.psAllMessages )
+    = Producer.plainSink( s.psAllMessages, Tracked.createProducer( s.psAllMessages ) )
       .contramap( (msg: Untracked[Value]) => s.record( msg.value ) )
 
 }

test/com/workflowfm/pew/stateless/KafkaExecutorTests.scala

@@ -1,8 +1,13 @@
 package com.workflowfm.pew.stateless
 
-import com.workflowfm.pew.stateless.StatelessMessages._
+import akka.Done
+import com.workflowfm.pew.stateless.StatelessMessages.{AnyMsg, _}
 import com.workflowfm.pew.stateless.instances.kafka.components.KafkaConnectors
-import com.workflowfm.pew._
+import com.workflowfm.pew.{PromiseHandler, _}
+import com.workflowfm.pew.stateless.components.{AtomicExecutor, Reducer, ResultListener}
+import com.workflowfm.pew.stateless.instances.kafka.MinimalKafkaExecutor
+import com.workflowfm.pew.stateless.instances.kafka.components.KafkaConnectors.{indyReducer, indySequencer, sendMessages}
+import org.apache.kafka.common.utils.Utils
 import org.bson.types.ObjectId
 import org.junit.runner.RunWith
 import org.scalatest.junit.JUnitRunner
@@ -17,6 +22,77 @@ class KafkaExecutorTests extends PewTestSuite with KafkaTests {
   // Ensure there are no outstanding messages before starting testing.
   new MessageDrain( true )
 
+  lazy val mainClassLoader: ClassLoader = Thread.currentThread().getContextClassLoader
+  lazy val kafkaClassLoader: ClassLoader = Utils.getContextOrKafkaClassLoader
+  lazy val threadClassLoader: ClassLoader = await( Future.unit.map(_ => Thread.currentThread().getContextClassLoader ) )
+
+  it should "use the same ClassLoader for Kafka as the Main thread" in {
+    mainClassLoader shouldBe kafkaClassLoader
+  }
+
+  // Instead unset the class loader when creating KafkaProducers.
+  ignore should "use the same ClassLoader for the ExecutionContext threads as the Main thread" in {
+    mainClassLoader shouldBe threadClassLoader
+  }
+
+  it should "start a producer from inside a ExecutionContext" in {
+    implicit val settings = completeProcess.settings
+
+    val future: Future[Done] = Future { Thread.sleep(100) } map { _ =>
+      val pii = PiInstance( ObjectId.get, pbi, PiObject(1) )
+      sendMessages( ReduceRequest( pii, Seq() ) )
+      Done
+    }
+
+    await( future ) shouldBe Done
+  }
+
+  it should "execute an atomic PbI using a DIY KafkaExecutor" in {
+    implicit val settings = completeProcess.settings
+    val listener = new ResultListener
+
+    val c1 = KafkaConnectors.indyReducer( new Reducer )
+    val c2 = KafkaConnectors.indySequencer
+    val c3 = KafkaConnectors.indyAtomicExecutor( AtomicExecutor() )
+    val c4 = KafkaConnectors.uniqueResultListener( listener )
+
+    val pii = PiInstance( ObjectId.get, pbi, PiObject(1) )
+    sendMessages( ReduceRequest( pii, Seq() ) )
+
+    val handler = new PromiseHandler( "test", pii.id )
+    listener.subscribe( handler )
+
+    await( handler.promise.future ) should be ("PbISleptFor1s")
+    await( KafkaConnectors.shutdown( c1, c2, c3, c4 ) )
+
+    val msgsOf = new MessageDrain( true )
+    msgsOf[SequenceRequest] shouldBe empty
+    msgsOf[SequenceFailure] shouldBe empty
+    msgsOf[ReduceRequest] shouldBe empty
+    msgsOf[Assignment] shouldBe empty
+    msgsOf[PiiUpdate] shouldBe empty
+  }
+
+  it should "execute an atomic PbI using the baremetal Executor interface" in {
+    val ex = makeExecutor( completeProcess.settings )
+
+    val piiId = await( ex.init( pbi, Seq( PiObject(1) ) ) )
+    val handler = new PromiseHandler( "test", piiId )
+    ex.subscribe( handler )
+
+    val f1 = handler.promise.future
+    ex.start( piiId )
+
+    await( f1 ) should be ("PbISleptFor1s")
+    ex.syncShutdown()
+
+    val msgsOf = new MessageDrain( true )
+    msgsOf[SequenceRequest] shouldBe empty
+    msgsOf[SequenceFailure] shouldBe empty
+    msgsOf[ReduceRequest] shouldBe empty
+    msgsOf[Assignment] shouldBe empty
+    msgsOf[PiiUpdate] shouldBe empty  }
+
   it should "execute atomic PbI once" in {
 
     val ex = makeExecutor( completeProcess.settings )