[fix][fn] Honor producer batching configuration in Python runtime

pulsar-functions/instance/src/main/python/contextimpl.py

@@ -172,17 +172,17 @@ def publish(self, topic_name, message, serde_class_name="serde.IdentitySerDe", p
     if compression_type is not None:
       pulsar_compression_type = compression_type
     if topic_name not in self.publish_producers:
+      producer_config = util.producer_config_from_function_details(self.instance_config.function_details)
       self.publish_producers[topic_name] = self.pulsar_client.create_producer(
         topic_name,
         block_if_queue_full=True,
-        batching_enabled=True,
-        batching_max_publish_delay_ms=10,
         compression_type=pulsar_compression_type,
         properties=util.get_properties(util.getFullyQualifiedFunctionName(
           self.instance_config.function_details.tenant,
           self.instance_config.function_details.namespace,
           self.instance_config.function_details.name),
-          self.instance_config.instance_id)
+          self.instance_config.instance_id),
+        **producer_config
       )
 
     if serde_class_name not in self.publish_serializers:

pulsar-functions/instance/src/main/python/python_instance.py

@@ -372,12 +372,7 @@ def setup_producer(self, producer_name=None):
             len(self.instance_config.function_details.sink.topic) > 0:
       Log.debug("Setting up producer for topic %s" % self.instance_config.function_details.sink.topic)
 
-      batch_type = pulsar.BatchingType.Default
-      if self.instance_config.function_details.sink.producerSpec.batchBuilder != None and \
-            len(self.instance_config.function_details.sink.producerSpec.batchBuilder) > 0:
-        batch_builder = self.instance_config.function_details.sink.producerSpec.batchBuilder
-        if batch_builder == "KEY_BASED":
-          batch_type = pulsar.BatchingType.KeyBased
+      producer_config = util.producer_config_from_function_details(self.instance_config.function_details)
 
       self.output_schema = self.get_schema(self.instance_config.function_details.sink.schemaType,
                                            self.instance_config.function_details.sink.typeClassName,
@@ -403,9 +398,6 @@ def setup_producer(self, producer_name=None):
         schema=self.output_schema,
         producer_name=producer_name,
         block_if_queue_full=True,
-        batching_enabled=True,
-        batching_type=batch_type,
-        batching_max_publish_delay_ms=10,
         compression_type=compression_type,
         # set send timeout to be infinity to prevent potential deadlock with consumer
         # that might happen when consumer is blocked due to unacked messages
@@ -417,7 +409,8 @@ def setup_producer(self, producer_name=None):
                         self.instance_config.function_details.tenant,
                         self.instance_config.function_details.namespace,
                         self.instance_config.function_details.name),
-                        self.instance_config.instance_id)
+                        self.instance_config.instance_id),
+        **producer_config
       )
 
   def setup_state(self):

pulsar-functions/instance/src/main/python/util.py

@@ -28,6 +28,7 @@
 import configparser
 
 from threading import Timer
+import pulsar
 from pulsar.functions import serde
 
 import log
@@ -82,6 +83,63 @@ def getFullyQualifiedInstanceId(tenant, namespace, name, instance_id):
 def get_properties(fullyQualifiedName, instanceId):
     return {"application": "pulsar-function", "id": str(fullyQualifiedName), "instance_id": str(instanceId)}
 
+# Keep these defaults aligned with ProducerBuilderFactory in the Java runtime.
+DEFAULT_BATCHING_ENABLED = True
+DEFAULT_BATCHING_MAX_PUBLISH_DELAY_MS = 10
+
+def batching_type_from_batch_builder(batch_builder):
+  if batch_builder == "KEY_BASED":
+    return pulsar.BatchingType.KeyBased
+  return pulsar.BatchingType.Default
+
+def producer_config_from_spec(producer_spec):
+  """Translate ProducerSpec batching and pending-queue fields to Python client arguments."""
+  config = {
+    "batching_enabled": DEFAULT_BATCHING_ENABLED,
+    "batching_max_publish_delay_ms": DEFAULT_BATCHING_MAX_PUBLISH_DELAY_MS,
+  }
+
+  if producer_spec is None:
+    return config
+
+  # ProducerSpec.batchBuilder predates BatchingSpec. Apply it first so the nested field can
+  # override it, matching ProducerBuilderFactory in the Java runtime.
+  if producer_spec.batchBuilder:
+    config["batching_type"] = batching_type_from_batch_builder(producer_spec.batchBuilder)
+
+  if producer_spec.maxPendingMessages > 0:
+    config["max_pending_messages"] = producer_spec.maxPendingMessages
+  if producer_spec.maxPendingMessagesAcrossPartitions > 0:
+    config["max_pending_messages_across_partitions"] = producer_spec.maxPendingMessagesAcrossPartitions
+
+  if not producer_spec.HasField("batchingSpec"):
+    return config
+
+  batching_spec = producer_spec.batchingSpec
+  # An empty nested proto is treated as an absent configuration by BatchingUtils.convertFromSpec.
+  # Public Function configs carry the 10ms default when enabled=false, so that case remains
+  # distinguishable from an empty message.
+  if not batching_spec.ListFields():
+    return config
+
+  config["batching_enabled"] = batching_spec.enabled
+  if batching_spec.batchingMaxPublishDelayMs > 0:
+    config["batching_max_publish_delay_ms"] = batching_spec.batchingMaxPublishDelayMs
+  if batching_spec.batchingMaxMessages > 0:
+    config["batching_max_messages"] = batching_spec.batchingMaxMessages
+  if batching_spec.batchingMaxBytes > 0:
+    config["batching_max_allowed_size_in_bytes"] = batching_spec.batchingMaxBytes
+  if batching_spec.batchBuilder:
+    config["batching_type"] = batching_type_from_batch_builder(batching_spec.batchBuilder)
+
+  # The Python client has no equivalent for roundRobinRouterBatchingPartitionSwitchFrequency.
+  return config
+
+def producer_config_from_function_details(function_details):
+  if function_details is None or not function_details.sink.HasField("producerSpec"):
+    return producer_config_from_spec(None)
+  return producer_config_from_spec(function_details.sink.producerSpec)
+
 def read_config(config_file):
     """
     The content of the configuration file is styled as follows:

pulsar-functions/instance/src/test/python/test_python_instance.py

@@ -37,7 +37,9 @@
 import Function_pb2
 import log
 import os
+import pulsar
 import unittest
+import util
 
 class TestContextImpl(unittest.TestCase):
 
@@ -149,3 +151,208 @@ def test_do_not_forward_properties(self):
     self.assertNotIn("custom-key", kwargs['properties'])
     self.assertIn("__pfn_input_topic__", kwargs['properties'])
 
+
+class TestProducerConfigFromSpec(unittest.TestCase):
+
+  def test_defaults_when_no_producer_spec(self):
+    function_details = Function_pb2.FunctionDetails()
+    self.assertEqual(util.producer_config_from_function_details(function_details), {
+      "batching_enabled": True,
+      "batching_max_publish_delay_ms": 10,
+    })
+
+  def test_empty_batching_spec_keeps_defaults(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchingSpec.SetInParent()
+    self.assertTrue(function_details.sink.producerSpec.HasField("batchingSpec"))
+    config = util.producer_config_from_function_details(function_details)
+    self.assertTrue(config["batching_enabled"])
+    self.assertEqual(config["batching_max_publish_delay_ms"], 10)
+
+  def test_batching_can_be_disabled(self):
+    function_details = Function_pb2.FunctionDetails()
+    batching_spec = function_details.sink.producerSpec.batchingSpec
+    batching_spec.enabled = False
+    batching_spec.batchingMaxPublishDelayMs = 10
+    config = util.producer_config_from_function_details(function_details)
+    self.assertFalse(config["batching_enabled"])
+    self.assertEqual(config["batching_max_publish_delay_ms"], 10)
+
+  def test_full_batching_spec_is_translated(self):
+    function_details = Function_pb2.FunctionDetails()
+    batching_spec = function_details.sink.producerSpec.batchingSpec
+    batching_spec.enabled = True
+    batching_spec.batchingMaxPublishDelayMs = 1
+    batching_spec.batchingMaxMessages = 500
+    batching_spec.batchingMaxBytes = 65536
+    batching_spec.batchBuilder = "KEY_BASED"
+    self.assertEqual(util.producer_config_from_function_details(function_details), {
+      "batching_enabled": True,
+      "batching_max_publish_delay_ms": 1,
+      "batching_max_messages": 500,
+      "batching_max_allowed_size_in_bytes": 65536,
+      "batching_type": pulsar.BatchingType.KeyBased,
+    })
+
+  def test_non_positive_values_keep_defaults(self):
+    function_details = Function_pb2.FunctionDetails()
+    batching_spec = function_details.sink.producerSpec.batchingSpec
+    batching_spec.enabled = True
+    batching_spec.batchingMaxPublishDelayMs = 0
+    batching_spec.batchingMaxMessages = 0
+    batching_spec.batchingMaxBytes = 0
+    config = util.producer_config_from_function_details(function_details)
+    self.assertEqual(config["batching_max_publish_delay_ms"], 10)
+    self.assertNotIn("batching_max_messages", config)
+    self.assertNotIn("batching_max_allowed_size_in_bytes", config)
+
+  def test_pending_message_limits_are_translated(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.maxPendingMessages = 2000
+    function_details.sink.producerSpec.maxPendingMessagesAcrossPartitions = 8000
+    config = util.producer_config_from_function_details(function_details)
+    self.assertEqual(config["max_pending_messages"], 2000)
+    self.assertEqual(config["max_pending_messages_across_partitions"], 8000)
+
+  def test_pending_message_limits_are_omitted_when_unset(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchingSpec.enabled = True
+    config = util.producer_config_from_function_details(function_details)
+    self.assertNotIn("max_pending_messages", config)
+    self.assertNotIn("max_pending_messages_across_partitions", config)
+
+  def test_producer_spec_batch_builder_is_honoured(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchBuilder = "KEY_BASED"
+    config = util.producer_config_from_function_details(function_details)
+    self.assertEqual(config["batching_type"], pulsar.BatchingType.KeyBased)
+
+  def test_batching_spec_batch_builder_overrides_producer_spec(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchBuilder = "KEY_BASED"
+    function_details.sink.producerSpec.batchingSpec.enabled = True
+    function_details.sink.producerSpec.batchingSpec.batchBuilder = "DEFAULT"
+    config = util.producer_config_from_function_details(function_details)
+    self.assertEqual(config["batching_type"], pulsar.BatchingType.Default)
+
+  def test_unknown_batch_builder_falls_back_to_default(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchBuilder = "UNKNOWN"
+    config = util.producer_config_from_function_details(function_details)
+    self.assertEqual(config["batching_type"], pulsar.BatchingType.Default)
+
+  def test_unsupported_round_robin_frequency_is_ignored(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchingSpec.enabled = True
+    function_details.sink.producerSpec.batchingSpec.roundRobinRouterBatchingPartitionSwitchFrequency = 7
+    config = util.producer_config_from_function_details(function_details)
+    self.assertNotIn("round_robin_router_batching_partition_switch_frequency", config)
+
+  def test_none_function_details_keeps_defaults(self):
+    self.assertEqual(util.producer_config_from_function_details(None), {
+      "batching_enabled": True,
+      "batching_max_publish_delay_ms": 10,
+    })
+
+
+class TestSinkProducerBatchingConfig(unittest.TestCase):
+
+  def _create_producer_kwargs(self, function_details):
+    pulsar_client = Mock()
+    pulsar_client.create_producer.return_value = Mock()
+    instance = PythonInstance('test_instance', 'test_func', '1.0', function_details, 100, 30,
+                              'user_code', pulsar_client, Mock(), 'test_cluster', 'test_url', None)
+    instance.get_schema = Mock(return_value="DEFAULT_SCHEMA")
+    instance.get_crypto_reader = Mock(return_value=None)
+    instance.setup_producer()
+    _, kwargs = pulsar_client.create_producer.call_args
+    return kwargs
+
+  def test_defaults_are_unchanged_without_a_producer_spec(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.topic = "test_sink_topic"
+    kwargs = self._create_producer_kwargs(function_details)
+    self.assertTrue(kwargs["batching_enabled"])
+    self.assertEqual(kwargs["batching_max_publish_delay_ms"], 10)
+    self.assertTrue(kwargs["block_if_queue_full"])
+
+  def test_batching_disabled_reaches_the_producer(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.topic = "test_sink_topic"
+    function_details.sink.producerSpec.batchingSpec.enabled = False
+    function_details.sink.producerSpec.batchingSpec.batchingMaxPublishDelayMs = 10
+    kwargs = self._create_producer_kwargs(function_details)
+    self.assertFalse(kwargs["batching_enabled"])
+
+  def test_batching_settings_reach_the_producer(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.topic = "test_sink_topic"
+    batching_spec = function_details.sink.producerSpec.batchingSpec
+    batching_spec.enabled = True
+    batching_spec.batchingMaxPublishDelayMs = 2
+    batching_spec.batchingMaxMessages = 100
+    batching_spec.batchingMaxBytes = 4096
+    function_details.sink.producerSpec.maxPendingMessages = 500
+    kwargs = self._create_producer_kwargs(function_details)
+    self.assertTrue(kwargs["batching_enabled"])
+    self.assertEqual(kwargs["batching_max_publish_delay_ms"], 2)
+    self.assertEqual(kwargs["batching_max_messages"], 100)
+    self.assertEqual(kwargs["batching_max_allowed_size_in_bytes"], 4096)
+    self.assertEqual(kwargs["max_pending_messages"], 500)
+
+  def test_key_based_batch_builder_reaches_the_producer(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.topic = "test_sink_topic"
+    function_details.sink.producerSpec.batchBuilder = "KEY_BASED"
+    kwargs = self._create_producer_kwargs(function_details)
+    self.assertEqual(kwargs["batching_type"], pulsar.BatchingType.KeyBased)
+
+
+class TestContextPublishBatchingConfig(unittest.TestCase):
+
+  def _create_producer_kwargs(self, function_details):
+    instance_config = InstanceConfig('test_instance_id', 'test_function_id', 'test_function_version',
+                                     function_details, 100)
+    pulsar_client = Mock()
+    producer = Mock()
+    producer.send_async = Mock(return_value=None)
+    pulsar_client.create_producer = Mock(return_value=producer)
+    context_impl = ContextImpl(instance_config, log.Log, pulsar_client, __file__, None, None, None, None, None)
+
+    msg = Message()
+    msg.message_id = Mock(return_value="test_message_id")
+    msg.partition_key = Mock(return_value="test_key")
+    context_impl.set_current_message_context(msg, "test_topic_name")
+    context_impl.publish("test_topic_name", "test_message")
+
+    _, kwargs = pulsar_client.create_producer.call_args
+    return kwargs
+
+  def test_defaults_are_unchanged_without_a_producer_spec(self):
+    kwargs = self._create_producer_kwargs(Function_pb2.FunctionDetails())
+    self.assertTrue(kwargs["batching_enabled"])
+    self.assertEqual(kwargs["batching_max_publish_delay_ms"], 10)
+    self.assertTrue(kwargs["block_if_queue_full"])
+
+  def test_batching_disabled_reaches_the_producer(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchingSpec.enabled = False
+    function_details.sink.producerSpec.batchingSpec.batchingMaxPublishDelayMs = 10
+    kwargs = self._create_producer_kwargs(function_details)
+    self.assertFalse(kwargs["batching_enabled"])
+
+  def test_batching_settings_reach_the_producer(self):
+    function_details = Function_pb2.FunctionDetails()
+    batching_spec = function_details.sink.producerSpec.batchingSpec
+    batching_spec.enabled = True
+    batching_spec.batchingMaxPublishDelayMs = 5
+    batching_spec.batchingMaxMessages = 250
+    kwargs = self._create_producer_kwargs(function_details)
+    self.assertEqual(kwargs["batching_max_publish_delay_ms"], 5)
+    self.assertEqual(kwargs["batching_max_messages"], 250)
+
+  def test_batch_builder_reaches_the_producer(self):
+    function_details = Function_pb2.FunctionDetails()
+    function_details.sink.producerSpec.batchBuilder = "KEY_BASED"
+    kwargs = self._create_producer_kwargs(function_details)
+    self.assertEqual(kwargs["batching_type"], pulsar.BatchingType.KeyBased)