[BEAM-13386] Add RLock support for cloudpickle (#16250)

ryanthompson591 · web-flow · commit fea9c4b4b1d1 · 2022-01-07T14:00:26.000-08:00
diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py
@@ -43,6 +43,7 @@
 # Pickling, especially unpickling, causes broken module imports on Python 3
 # if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884.
 _pickle_lock = threading.RLock()
+RLOCK_TYPE = type(_pickle_lock)
 
 
 def dumps(o, enable_trace=True, use_zlib=False):
@@ -56,6 +57,10 @@ def dumps(o, enable_trace=True, use_zlib=False):
         pickler.dispatch_table[type(flags.FLAGS)] = _pickle_absl_flags
       except NameError:
         pass
+      try:
+        pickler.dispatch_table[RLOCK_TYPE] = _pickle_rlock
+      except NameError:
+        pass
       pickler.dump(o)
       s = file.getvalue()
 
@@ -99,6 +104,10 @@ def _create_absl_flags():
   return flags.FLAGS
 
 
+def _pickle_rlock(obj):
+  return RLOCK_TYPE, tuple([])
+
+
 def dump_session(file_path):
   # It is possible to dump session with cloudpickle. However, since references
   # are saved it should not be necessary. See https://s.apache.org/beam-picklers
diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py b/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py
@@ -20,6 +20,7 @@
 # pytype: skip-file
 
 import sys
+import threading
 import types
 import unittest
 
@@ -91,6 +92,12 @@ def foo():
 
     self.assertEqual('expected_value', loads(dumps(foo))())
 
+  def test_pickle_rlock(self):
+    rlock_instance = threading.RLock()
+    rlock_type = type(rlock_instance)
+
+    self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type)
+
   @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced')
   def test_dump_and_load_mapping_proxy(self):
     self.assertEqual(
diff --git a/sdks/python/apache_beam/internal/pickler_test.py b/sdks/python/apache_beam/internal/pickler_test.py
@@ -20,6 +20,7 @@
 # pytype: skip-file
 
 import sys
+import threading
 import types
 import unittest
 
@@ -87,6 +88,12 @@ def test_recursive_class(self):
         'RecursiveClass:abc',
         loads(dumps(module_test.RecursiveClass('abc').datum)))
 
+  def test_pickle_rlock(self):
+    rlock_instance = threading.RLock()
+    rlock_type = type(rlock_instance)
+
+    self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type)
+
   @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced')
   def test_dump_and_load_mapping_proxy(self):
     self.assertEqual(
