def _open(self, **kwargs):
self._pos = 0
if self._buf is None:
if self._mode == OpenMode.OPEN_WRITE:
self._written = 0
self._buf = DlgSharedMemory(self._name)
return self._buf
def test_open_twice(self):
"""
It should be possible to open a shared memory block simultaneously and therefore offer a
view into the same memory buffer.
This does bring to light that this primitive is not thread-safe by default.
"""
block_a = DlgSharedMemory("A")
block_b = DlgSharedMemory("A")
self.assertEqual(block_a.buf, block_b.buf)
block_a.close()
block_b.close()
block_a.unlink()
def test_resize_shrink(self):
"""
Size should shrink when requested
"""
block_a = DlgSharedMemory("A")
old_size = block_a.size
block_a.resize(old_size // 2)
self.assertEqual(block_a.size, old_size // 2)
block_a.close()
block_a.unlink()
def test_resize_grow(self):
"""
The size should grow when requested
"""
block_a = DlgSharedMemory("A")
old_size = block_a.size
block_a.resize(block_a.size * 2)
self.assertEqual(block_a.size, old_size * 2)
block_a.close()
block_a.unlink()
def test_open_random(self):
"""
It should be possible to open a randomly named block (although this must be deliberate)
"""
block_a = DlgSharedMemory(None)
self.assertIsNotNone(block_a.name)
block_a.close()
block_a.unlink()
def test_open_new(self):
"""
When opening a new block by name, that name should be respected
"""
block_a = DlgSharedMemory("A")
self.assertEqual(block_a.name, "/A")
block_a.close()
block_a.unlink()
def test_long_name(self):
"""
There is a maximum name size which should be respected
"""
filename = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
block_a = DlgSharedMemory(filename)
self.assertEqual(block_a.name, "/" + filename[:_MAXNAMELENGTH - 1])
block_a.close()
block_a.unlink()
def test_start_slash(self):
"""
If a user attempts to make a block with a name beginning with /, it should be ignored
(rather than having // as a name)
"""
filename = "/memory"
block_a = DlgSharedMemory(filename)
self.assertEqual(block_a.name, filename)
block_a.close()
block_a.unlink()
def _close_block(session_id, name):
mem = DlgSharedMemory(f"{session_id}_{name}")
mem.close()
def _cleanup_block(session_id, name):
mem = DlgSharedMemory(f"{session_id}_{name}")
mem.close()
mem.unlink(
) # It is unlinking that is critical to freeing resources from the OS
def test_open_existing(self):
"""
When opening an existing but still linked block, data should persist
"""
block_a = DlgSharedMemory("A")
data = pickle.dumps(3)
block_a.buf[0:len(data)] = data
block_a.resize(len(data))
old_size = block_a.size
block_a.close()
block_b = DlgSharedMemory("A")
self.assertEqual(block_b.size, old_size)
self.assertEqual(block_b.name, "/A")
self.assertEqual(block_b.buf[:], data)
block_b.close()
block_b.unlink()
def test_unlink_twice(self):
"""
It should not be possible to unlink a block twice.
"""
block_a = DlgSharedMemory("A")
block_b = DlgSharedMemory("A")
block_a.close()
block_b.close()
block_a.unlink()
with self.assertWarns(RuntimeWarning):
block_b.unlink()
class SharedMemoryIO(DataIO):
"""
A DataIO class that writes to a shared memory buffer
"""
def __init__(self, uid, session_id, **kwargs):
super().__init__()
self._name = f"{session_id}_{uid}"
self._written = 0
self._pos = 0
self._buf = None
@overrides
def _open(self, **kwargs):
self._pos = 0
if self._buf is None:
if self._mode == OpenMode.OPEN_WRITE:
self._written = 0
self._buf = DlgSharedMemory(self._name)
return self._buf
@overrides
def _write(self, data, **kwargs) -> int:
total_size = len(data) + self._written
if total_size > self._buf.size:
self._buf.resize(total_size)
self._buf.buf[self._written:total_size] = data
self._written = total_size
else:
self._buf.buf[self._written:total_size] = data
self._written = total_size
self._buf.resize(total_size)
# It may be inefficient to resize many times, but assuming data is written 'once' this is
# might be tolerable and guarantees that the size of the underlying buffer is tight.
return len(data)
@overrides
def _read(self, count=4096, **kwargs):
if self._pos == self._buf.size:
return None
start = self._pos
end = self._pos + count
end = min(end, self._buf.size)
out = self._buf.buf[start:end]
self._pos = end
return out
@overrides
def _close(self, **kwargs):
if self._mode == OpenMode.OPEN_WRITE:
self._buf.resize(self._written)
self._buf.close()
self._buf = None
@overrides
def _size(self, **kwargs) -> int:
return self._buf.size
@overrides
def exists(self) -> bool:
return self._buf is not None
@overrides
def delete(self):
self._close()