diff --git a/flink-python/pyflink/datastream/state.py b/flink-python/pyflink/datastream/state.py index 6659625e12c5a0..65fcb9757d2e97 100644 --- a/flink-python/pyflink/datastream/state.py +++ b/flink-python/pyflink/datastream/state.py @@ -834,7 +834,7 @@ def cleanup_in_rocksdb_compact_filter( StateTtlConfig.CleanupStrategies.Strategies.ROCKSDB_COMPACTION_FILTER] = \ StateTtlConfig.CleanupStrategies.RocksdbCompactFilterCleanupStrategy( query_time_after_num_entries, - periodic_compaction_time if periodic_compaction_time else Duration.of_days(30)) + periodic_compaction_time) return self def disable_cleanup_in_background(self) -> 'StateTtlConfig.Builder': @@ -928,13 +928,15 @@ def __init__(self, query_time_after_num_entries: int, periodic_compaction_time=None): self._query_time_after_num_entries = query_time_after_num_entries - self._periodic_compaction_time = periodic_compaction_time \ - if periodic_compaction_time else Duration.of_days(30) + # Creating the default Duration here would start a Py4J gateway in embedded workers. + self._periodic_compaction_time = periodic_compaction_time or None def get_query_time_after_num_entries(self) -> int: return self._query_time_after_num_entries def get_periodic_compaction_time(self) -> Duration: + if self._periodic_compaction_time is None: + return Duration.of_days(30) return self._periodic_compaction_time EMPTY_STRATEGY = EmptyCleanupStrategy() diff --git a/flink-python/pyflink/datastream/tests/test_data_stream.py b/flink-python/pyflink/datastream/tests/test_data_stream.py index 527439afa3ea62..272473b8450475 100644 --- a/flink-python/pyflink/datastream/tests/test_data_stream.py +++ b/flink-python/pyflink/datastream/tests/test_data_stream.py @@ -1255,6 +1255,98 @@ def setUp(self): config = get_j_env_configuration(self.env._j_stream_execution_environment) config.setString("python.execution-mode", "thread") + def test_gateway_access_is_rejected(self): + self.addCleanup(self.env.set_parallelism, self.env.get_parallelism()) + self.env.set_parallelism(1) + + class GatewayAccessFunction(MapFunction): + def open(self, runtime_context: RuntimeContext): + from pemja import findClass + + self.j_integer = findClass('java.lang.Integer') + + def map(self, value): + from unittest import TestCase + from unittest.mock import patch + + from pyflink.java_gateway import get_gateway, launch_gateway + + with patch('pyflink.java_gateway.launch_gateway_server_process', + side_effect=AssertionError('Gateway subprocess launched')): + for gateway_function in [get_gateway, launch_gateway]: + with TestCase().assertRaisesRegex(RuntimeError, r'pemja\.findClass'): + gateway_function() + + return self.j_integer.parseInt(value) + + (self.env.from_collection(['1', '2'], type_info=Types.STRING()) + .map(GatewayAccessFunction(), output_type=Types.INT()) + .add_sink(self.test_sink)) + self.env.execute('test_gateway_access_is_rejected') + self.assert_equals_sorted(['1', '2'], self.test_sink.get_results()) + + def test_state_ttl_without_gateway(self): + self.addCleanup(self.env.set_parallelism, self.env.get_parallelism()) + self.env.set_parallelism(1) + config = get_j_env_configuration(self.env._j_stream_execution_environment) + if config.containsKey("state.backend.type"): + self.addCleanup(config.setString, "state.backend.type", + config.getString("state.backend.type", None)) + else: + self.addCleanup(config.removeKey, "state.backend.type") + config.setString("state.backend.type", "rocksdb") + + class TtlStateFunction(KeyedProcessFunction): + def open(self, runtime_context: RuntimeContext): + from unittest.mock import patch + + from pyflink.fn_execution.embedded.java_utils import to_java_state_ttl_config + + self.states = [] + cases = [('default', None, True), + ('explicit_default', 17, True), + ('disabled', None, False), + ('explicit_disabled', 23, False)] + # Guard the actual embedded worker, after the client gateway already exists. + with patch('pyflink.common.time.get_gateway', side_effect=AssertionError( + 'TTL state initialization must not request a Py4J gateway')): + for name, queries, background in cases: + builder = StateTtlConfig.new_builder(Time.days(1)) + if queries is not None: + builder.cleanup_in_rocksdb_compact_filter(queries) + if not background: + builder.disable_cleanup_in_background() + ttl_config = builder.build() + descriptor = ValueStateDescriptor('ttl_' + name, Types.INT()) + descriptor.enable_time_to_live(ttl_config) + self.states.append(runtime_context.get_state(descriptor)) + + j_ttl_config = to_java_state_ttl_config(ttl_config) + assert j_ttl_config.getTimeToLive().toMillis() == 86400000 + cleanup = j_ttl_config.getCleanupStrategies() + assert cleanup.isCleanupInBackground() == background + rocksdb_cleanup = cleanup.getRocksdbCompactFilterCleanupStrategy() + if background or queries is not None: + assert rocksdb_cleanup.getQueryTimeAfterNumEntries() == ( + 1000 if queries is None else queries) + assert rocksdb_cleanup.getPeriodicCompactionTime().toMillis() == ( + 30 * 24 * 60 * 60 * 1000) + else: + assert rocksdb_cleanup is None + + def process_element(self, value, ctx): + for index, state in enumerate(self.states): + state.update((state.value() or 0) + value) + yield index, state.value() + + (self.env.from_collection([1, 2], type_info=Types.INT()) + .key_by(lambda value: 0) + .process(TtlStateFunction(), output_type=Types.TUPLE([Types.INT(), Types.INT()])) + .add_sink(self.test_sink)) + self.env.execute('test_state_ttl_without_gateway') + expected = ['(%s,%s)' % (index, total) for index in range(4) for total in [1, 3]] + self.assert_equals_sorted(expected, self.test_sink.get_results()) + def test_metrics(self): ds = self.env.from_collection( [('ab', 'a', decimal.Decimal(1)), diff --git a/flink-python/pyflink/datastream/tests/test_state_ttl_config.py b/flink-python/pyflink/datastream/tests/test_state_ttl_config.py new file mode 100644 index 00000000000000..c74a27c900813e --- /dev/null +++ b/flink-python/pyflink/datastream/tests/test_state_ttl_config.py @@ -0,0 +1,91 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ +import pickle +import unittest +from unittest.mock import patch + +from pyflink.common.time import Duration, Time +from pyflink.datastream.state import StateTtlConfig + + +class StateTtlConfigTests(unittest.TestCase): + + def test_default_cleanup_without_gateway(self): + with patch('pyflink.common.time.get_gateway', + side_effect=AssertionError('Gateway requested')): + config = StateTtlConfig.new_builder(Time.days(1)).build() + strategy = config.get_cleanup_strategies().get_rocksdb_compact_filter_cleanup_strategy() + self.assertEqual(strategy.get_query_time_after_num_entries(), 1000) + + def test_explicit_cleanup_without_gateway(self): + for arguments in [(), (None,), (0,), (False,), ('',)]: + with self.subTest(arguments=arguments): + with patch('pyflink.common.time.get_gateway', + side_effect=AssertionError('Gateway requested')): + config = (StateTtlConfig.new_builder(Time.days(1)) + .cleanup_in_rocksdb_compact_filter(17, *arguments) + .build()) + strategy = (config.get_cleanup_strategies() + .get_rocksdb_compact_filter_cleanup_strategy()) + self.assertEqual(strategy.get_query_time_after_num_entries(), 17) + + def test_cleanup_round_trip_without_gateway(self): + for explicit_cleanup in [False, True]: + for serializer in ['pickle', 'protobuf']: + with self.subTest(explicit_cleanup=explicit_cleanup, serializer=serializer): + with patch('pyflink.common.time.get_gateway', + side_effect=AssertionError('Gateway requested')): + builder = StateTtlConfig.new_builder(Time.days(1)) + if explicit_cleanup: + builder.cleanup_in_rocksdb_compact_filter(17) + config = builder.build() + if serializer == 'pickle': + restored = pickle.loads(pickle.dumps(config)) + else: + restored = StateTtlConfig._from_proto(config._to_proto()) + strategy = (restored.get_cleanup_strategies() + .get_rocksdb_compact_filter_cleanup_strategy()) + self.assertEqual(restored.get_ttl(), Time.days(1)) + self.assertEqual(strategy.get_query_time_after_num_entries(), + 17 if explicit_cleanup else 1000) + + def test_default_duration_is_created_by_accessor(self): + with patch('pyflink.common.time.get_gateway') as gateway: + strategy = StateTtlConfig.CleanupStrategies.RocksdbCompactFilterCleanupStrategy(17) + gateway.assert_not_called() + + duration = strategy.get_periodic_compaction_time() + + gateway.assert_called_once_with() + java_duration = gateway.return_value.jvm.java.time.Duration + java_duration.ofDays.assert_called_once_with(30) + self.assertIs(duration._j_duration, java_duration.ofDays.return_value) + + def test_explicit_duration_is_preserved(self): + for milliseconds in [0, 3600000]: + with self.subTest(milliseconds=milliseconds): + with patch('pyflink.common.time.get_gateway'): + duration = Duration.of_millis(milliseconds) + with patch('pyflink.common.time.get_gateway', + side_effect=AssertionError('Gateway requested')): + config = (StateTtlConfig.new_builder(Time.days(1)) + .cleanup_in_rocksdb_compact_filter(17, duration) + .build()) + strategy = (config.get_cleanup_strategies() + .get_rocksdb_compact_filter_cleanup_strategy()) + self.assertIs(strategy.get_periodic_compaction_time(), duration) diff --git a/flink-python/pyflink/fn_execution/tests/test_java_gateway.py b/flink-python/pyflink/fn_execution/tests/test_java_gateway.py new file mode 100644 index 00000000000000..3e66af00f3947b --- /dev/null +++ b/flink-python/pyflink/fn_execution/tests/test_java_gateway.py @@ -0,0 +1,75 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ +import os +import sys +import unittest +from unittest.mock import patch + +from pyflink import java_gateway + + +class JavaGatewayTests(unittest.TestCase): + + def test_embedded_gateway_access_is_rejected_before_locking(self): + for cached_gateway in [None, object()]: + for gateway_port in [None, '12345']: + with self.subTest(cached=cached_gateway is not None, port=gateway_port): + with patch.dict(sys.modules, {'_pemja': object()}), \ + patch.dict(os.environ), \ + patch.object(java_gateway, '_gateway', cached_gateway), \ + patch.object(java_gateway, '_lock') as lock: + os.environ.pop('PYFLINK_GATEWAY_PORT', None) + if gateway_port is not None: + os.environ['PYFLINK_GATEWAY_PORT'] = gateway_port + lock.__enter__.side_effect = AssertionError('Gateway lock acquired') + + with self.assertRaisesRegex(RuntimeError, r'pemja\.findClass'): + java_gateway.get_gateway() + + lock.__enter__.assert_not_called() + + def test_embedded_gateway_launch_is_rejected(self): + with patch.dict(sys.modules, {'_pemja': object()}), \ + patch.object(java_gateway, 'is_launch_gateway_disabled', return_value=False), \ + patch.object(java_gateway, '_find_flink_home'), \ + patch.object(java_gateway, 'launch_gateway_server_process') as launch: + launch.side_effect = AssertionError('Gateway subprocess launched') + + with self.assertRaisesRegex(RuntimeError, r'pemja\.findClass'): + java_gateway.launch_gateway() + + launch.assert_not_called() + + def test_client_gateway_is_available_when_pemja_package_is_imported(self): + cached_gateway = object() + with patch.dict(sys.modules, {'pemja': object()}), \ + patch.object(java_gateway, '_gateway', cached_gateway): + sys.modules.pop('_pemja', None) + + self.assertIs(java_gateway.get_gateway(), cached_gateway) + + def test_process_worker_gateway_launch_is_still_disabled(self): + with patch.dict(sys.modules), \ + patch.dict(os.environ, {'PYFLINK_GATEWAY_DISABLED': 'true'}), \ + patch.object(java_gateway, 'launch_gateway_server_process') as launch: + sys.modules.pop('_pemja', None) + + with self.assertRaisesRegex(Exception, 'during Python UDF execution'): + java_gateway.launch_gateway() + + launch.assert_not_called() diff --git a/flink-python/pyflink/java_gateway.py b/flink-python/pyflink/java_gateway.py index e73ce51b0d8f5b..c019421b65f541 100644 --- a/flink-python/pyflink/java_gateway.py +++ b/flink-python/pyflink/java_gateway.py @@ -20,6 +20,7 @@ import shlex import shutil import struct +import sys import tempfile import time from logging import WARN @@ -44,8 +45,17 @@ def is_launch_gateway_disabled(): return False +def _check_gateway_access(): + # Pemja registers its native module before loading code in an embedded interpreter. + if '_pemja' in sys.modules: + raise RuntimeError( + 'Py4J gateway access is not supported in Python thread mode. ' + 'Use pemja.findClass to access Java classes.') + + def get_gateway() -> JavaGateway: global _gateway + _check_gateway_access() with _lock: if _gateway is None: # Set the level to WARN to mute the noisy INFO level logs @@ -81,6 +91,7 @@ def launch_gateway() -> JavaGateway: """ launch jvm gateway """ + _check_gateway_access() if is_launch_gateway_disabled(): raise Exception("It's launching the PythonGatewayServer during Python UDF execution " "which is unexpected. It usually happens when the job codes are "