Test Pickle with out-of-band buffers

distributed/protocol/tests/test_pickle.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from distributed.protocol.pickle import dumps, loads
+from distributed.protocol.pickle import HIGHEST_PROTOCOL, dumps, loads
 
 
 def test_pickle_data():
@@ -15,6 +15,37 @@ def test_pickle_data():
         assert loads(dumps(d)) == d
 
 
+def test_pickle_out_of_band():
+    try:
+        from pickle import PickleBuffer
+    except ImportError:
+        pass
+
+    class MemoryviewHolder:
+        def __init__(self, mv):
+            self.mv = memoryview(mv)
+
+        def __reduce_ex__(self, protocol):
+            if protocol >= 5:
+                return MemoryviewHolder, (PickleBuffer(self.mv),)
+            else:
+                return MemoryviewHolder, (self.mv.tobytes(),)
+
+    mv = memoryview(b"123")
+    mvh = MemoryviewHolder(mv)
+
+    if HIGHEST_PROTOCOL >= 5:
+        l = []
+        d = dumps(mvh, buffer_callback=l.append)
+        mvh2 = loads(d, buffers=l)
+    else:
+        mvh2 = loads(dumps(mvh))
+
+    assert isinstance(mvh2, MemoryviewHolder)
+    assert isinstance(mvh2.mv, memoryview)
+    assert mvh2.mv == mv
+
+
 def test_pickle_numpy():
     np = pytest.importorskip("numpy")
     x = np.ones(5)
@@ -23,6 +54,12 @@ def test_pickle_numpy():
     x = np.ones(5000)
     assert (loads(dumps(x)) == x).all()
 
+    if HIGHEST_PROTOCOL >= 5:
+        x = np.ones(5000)
+        l = []
+        d = dumps(x, buffer_callback=l.append)
+        assert (loads(d, buffers=l) == x).all()
+
 
 @pytest.mark.xfail(
     sys.version_info[:2] == (3, 8),