Don't deep-copy read-only buffers on unpickle (#8609)

Author: crusaderky

distributed/protocol/serialize.py

@@ -6,6 +6,7 @@
 from array import array
 from enum import Enum
 from functools import partial
+from pickle import PickleBuffer
 from types import ModuleType
 from typing import Any, Generic, Literal, TypeVar
 
@@ -86,19 +87,27 @@ def buffer_callback(f):
     return header, frames
 
 
-def pickle_loads(header, frames):
-    x, buffers = frames[0], frames[1:]
+def pickle_loads(
+    header: dict[str, Any], frames: list[bytes | bytearray | memoryview | PickleBuffer]
+) -> Any:
+    pik, buffers = frames[0], frames[1:]
 
-    writeable = header.get("writeable")
-    if not writeable:
-        writeable = len(buffers) * (None,)
+    def ensure_writeable_flag(mv: memoryview, w: bool) -> memoryview:
+        if w and mv.readonly:
+            # Can't avoid a deep copy
+            return memoryview(bytearray(mv))
+        elif not w and not mv.readonly:
+            # Zero copy - this is just a flag
+            return mv.toreadonly()
+        else:
+            return mv
 
     buffers = [
-        memoryview(bytearray(mv) if w else bytes(mv)) if w == mv.readonly else mv
-        for w, mv in zip(writeable, map(ensure_memoryview, buffers))
+        ensure_writeable_flag(ensure_memoryview(mv), w)
+        for mv, w in zip(buffers, header["writeable"])
     ]
 
-    return pickle.loads(x, buffers=buffers)
+    return pickle.loads(pik, buffers=buffers)
 
 
 def import_allowed_module(name):

distributed/protocol/tests/test_pickle.py

@@ -297,3 +297,21 @@ def func(df):
         out, _ = proc.communicate(timeout=60)
 
     assert "script successful" in out.decode("utf-8")
+
+
+@pytest.mark.parametrize("serializer", ["dask", "pickle"])
+def test_pickle_zero_copy_read_only_flag(serializer):
+    np = pytest.importorskip("numpy")
+    a = np.arange(10)
+    a.flags.writeable = False
+    header, frames = serialize(a, serializers=[serializer])
+    frames = [bytearray(f) for f in frames]  # Simulate network transfer
+    b = deserialize(header, frames)
+    c = deserialize(header, frames)
+    assert not b.flags.writeable
+    assert not c.flags.writeable
+    ptr_a = a.__array_interface__["data"][0]
+    ptr_b = b.__array_interface__["data"][0]
+    ptr_c = c.__array_interface__["data"][0]
+    assert ptr_b != ptr_a
+    assert ptr_b == ptr_c