From c0d6a5cea1bb35f3f1f9a47fdf58f5560c5ce258 Mon Sep 17 00:00:00 2001 From: Rui Fu Date: Thu, 20 Aug 2026 22:37:48 +0800 Subject: [PATCH] [fix][fn] Honor producer batching configuration in Python runtime --- .../instance/src/main/python/contextimpl.py | 6 +- .../src/main/python/python_instance.py | 13 +- .../instance/src/main/python/util.py | 58 +++++ .../src/test/python/test_python_instance.py | 207 ++++++++++++++++++ 4 files changed, 271 insertions(+), 13 deletions(-) diff --git a/pulsar-functions/instance/src/main/python/contextimpl.py b/pulsar-functions/instance/src/main/python/contextimpl.py index 826ad65b2e024..117cded4323b8 100755 --- a/pulsar-functions/instance/src/main/python/contextimpl.py +++ b/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: diff --git a/pulsar-functions/instance/src/main/python/python_instance.py b/pulsar-functions/instance/src/main/python/python_instance.py index 5c57dfef79008..a6b98cfd26869 100755 --- a/pulsar-functions/instance/src/main/python/python_instance.py +++ b/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): diff --git a/pulsar-functions/instance/src/main/python/util.py b/pulsar-functions/instance/src/main/python/util.py index 7b2a2d7b4d172..d79a137994b54 100755 --- a/pulsar-functions/instance/src/main/python/util.py +++ b/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: diff --git a/pulsar-functions/instance/src/test/python/test_python_instance.py b/pulsar-functions/instance/src/test/python/test_python_instance.py index 1e72db8545816..ae125c786e3e6 100644 --- a/pulsar-functions/instance/src/test/python/test_python_instance.py +++ b/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)