Changed: Change partition recovery hierachy in RdKafkaProducer

Author: eudiegoborgs

pkg/rdkafka/RdKafkaProducer.php

@@ -28,6 +28,22 @@ public function __construct(VendorProducer $producer, Serializer $serializer)
         $this->setSerializer($serializer);
     }
 
+
+    /**
+     * @param RdKafkaTopic   $destination
+     * @param RdKafkaMessage $message
+     */
+    private function getPartition(Destination $destination, Message $message): int
+    {
+        if (!is_null($message->getPartition())) {
+            return $message->getPartition();
+        }
+        if (!is_null($destination->getPartition())) {
+            return $destination->getPartition();
+        }
+        return RD_KAFKA_PARTITION_UA;
+    }
+
     /**
      * @param RdKafkaTopic   $destination
      * @param RdKafkaMessage $message
@@ -37,7 +53,7 @@ public function send(Destination $destination, Message $message): void
         InvalidDestinationException::assertDestinationInstanceOf($destination, RdKafkaTopic::class);
         InvalidMessageException::assertMessageInstanceOf($message, RdKafkaMessage::class);
 
-        $partition = $message->getPartition() ?: $destination->getPartition() ?: RD_KAFKA_PARTITION_UA;
+        $partition = $this->getPartition($destination, $message);
         $payload = $this->serializer->toString($message);
         $key = $message->getKey() ?: $destination->getKey() ?: null;
 

pkg/rdkafka/Tests/RdKafkaProducerTest.php

@@ -225,6 +225,104 @@ public function testShouldAllowSerializersToSerializeKeys()
         $producer->send(new RdKafkaTopic('theQueueName'), $message);
     }
 
+    public function testShouldGetPartitionFromMessage(): void
+    {
+        $partition = 1;
+
+        $kafkaTopic = $this->createKafkaTopicMock();
+        $kafkaTopic
+            ->expects($this->once())
+            ->method('producev')
+            ->with(
+                $partition,
+                0,
+                'theSerializedMessage',
+                'theSerializedKey'
+            )
+        ;
+
+        $kafkaProducer = $this->createKafkaProducerMock();
+        $kafkaProducer
+            ->expects($this->once())
+            ->method('newTopic')
+            ->willReturn($kafkaTopic)
+        ;
+        $kafkaProducer
+            ->expects($this->once())
+            ->method('poll')
+            ->with(0)
+        ;
+        $messageHeaders = ['bar' => 'barVal'];
+        $message = new RdKafkaMessage('theBody', ['foo' => 'fooVal'], $messageHeaders);
+        $message->setKey('key');
+        $message->setPartition($partition);
+
+        $serializer = $this->createSerializerMock();
+        $serializer
+            ->expects($this->once())
+            ->method('toString')
+            ->willReturnCallback(function () use ($message) {
+                $message->setKey('theSerializedKey');
+
+                return 'theSerializedMessage';
+            })
+        ;
+
+        $destination = new RdKafkaTopic('theQueueName');
+
+        $producer = new RdKafkaProducer($kafkaProducer, $serializer);
+        $producer->send($destination, $message);
+    }
+
+    public function testShouldGetPartitionFromDestination(): void
+    {
+        $partition = 2;
+
+        $kafkaTopic = $this->createKafkaTopicMock();
+        $kafkaTopic
+            ->expects($this->once())
+            ->method('producev')
+            ->with(
+                $partition,
+                0,
+                'theSerializedMessage',
+                'theSerializedKey'
+            )
+        ;
+
+        $kafkaProducer = $this->createKafkaProducerMock();
+        $kafkaProducer
+            ->expects($this->once())
+            ->method('newTopic')
+            ->willReturn($kafkaTopic)
+        ;
+        $kafkaProducer
+            ->expects($this->once())
+            ->method('poll')
+            ->with(0)
+        ;
+        $messageHeaders = ['bar' => 'barVal'];
+        $message = new RdKafkaMessage('theBody', ['foo' => 'fooVal'], $messageHeaders);
+        $message->setKey('key');
+
+        $serializer = $this->createSerializerMock();
+        $serializer
+            ->expects($this->once())
+            ->method('toString')
+            ->willReturnCallback(function () use ($message) {
+                $message->setKey('theSerializedKey');
+
+                return 'theSerializedMessage';
+            })
+        ;
+
+        $destination = new RdKafkaTopic('theQueueName');
+        $destination->setPartition($partition);
+
+        $producer = new RdKafkaProducer($kafkaProducer, $serializer);
+        $producer->send($destination, $message);
+    }
+
     /**
      * @return \PHPUnit\Framework\MockObject\MockObject|ProducerTopic
      */