Merge branch 'enh-tests-neilsimplify-shmem' into enh-tests-shmem

applio · applio · commit 8cf9ba3ac58d · 2019-02-16T19:27:54.000-06:00
adopting Neil's good work centralizing name autogeneration and removing
custom exceptions.
diff --git a/Lib/multiprocessing/shared_memory.py b/Lib/multiprocessing/shared_memory.py
@@ -1,4 +1,8 @@
-"Provides shared memory for direct access across processes."
+"""Provides shared memory for direct access across processes.
+
+The API of this package is currently provisional. Refer to the
+documentation for details.
+"""
 
 
 __all__ = [ 'SharedMemory', 'PosixSharedMemory', 'WindowsNamedSharedMemory',
@@ -12,23 +16,17 @@
                       BarrierProxy, AcquirerProxy, ConditionProxy, EventProxy
 from . import util
 import os
-import random
 import struct
 import sys
 import threading
+import secrets
 try:
-    from _posixshmem import _PosixSharedMemory, \
-                            Error, ExistentialError, PermissionsError, \
-                            O_CREAT, O_EXCL, O_CREX, O_TRUNC
+    import  _posixshmem
+    from os import O_CREAT, O_EXCL, O_TRUNC
 except ImportError as ie:
-    if os.name != "nt":
-        # On Windows, posixshmem is not required to be available.
-        raise ie
-    else:
-        _PosixSharedMemory = object
-        class Error(Exception): pass
-        class ExistentialError(Error): pass
-        O_CREAT, O_EXCL, O_CREX, O_TRUNC = 64, 128, 192, 512
+    O_CREAT, O_EXCL, O_TRUNC = 64, 128, 512
+
+O_CREX = O_CREAT | O_EXCL
 
 if os.name == "nt":
     import ctypes
@@ -89,15 +87,15 @@ class WindowsNamedSharedMemory:
 
     def __init__(self, name, flags=None, mode=384, size=0, read_only=False):
         if name is None:
-            name = f'wnsm_{os.getpid()}_{random.randrange(100000)}'
+            name = _make_filename()
 
         if size == 0:
             # Attempt to dynamically determine the existing named shared
             # memory block's size which is likely a multiple of mmap.PAGESIZE.
             try:
                 h_map = kernel32.OpenFileMappingW(FILE_MAP_READ, False, name)
-            except OSError as ose:
-                raise ExistentialError(*ose.args)
+            except OSError:
+                raise FileNotFoundError(name)
             try:
                 p_buf = kernel32.MapViewOfFile(h_map, FILE_MAP_READ, 0, 0, 0)
             finally:
@@ -115,57 +113,146 @@ def __init__(self, name, flags=None, mode=384, size=0, read_only=False):
             except OSError as ose:
                 name_collision = False
             if name_collision:
-                raise ExistentialError(
-                    f"Shared memory already exists with name={name}"
-                )
+                raise FileExistsError(name)
 
         self._mmap = mmap.mmap(-1, size, tagname=name)
-        self.buf = memoryview(self._mmap)
+        self._buf = memoryview(self._mmap)
         self.name = name
         self.mode = mode
-        self.size = size
+        self._size = size
+
+    @property
+    def size(self):
+        "Size in bytes."
+        return self._size
+
+    @property
+    def buf(self):
+        "A memoryview of contents of the shared memory block."
+        return self._buf
 
     def __repr__(self):
         return f'{self.__class__.__name__}({self.name!r}, size={self.size})'
 
     def close(self):
-        self.buf.release()
-        self._mmap.close()
+        if self._buf is not None:
+            self._buf.release()
+            self._buf = None
+        if self._mmap is not None:
+            self._mmap.close()
+            self._mmap = None
 
     def unlink(self):
         """Windows ensures that destruction of the last reference to this
         named shared memory block will result in the release of this memory."""
         pass
 
 
-class PosixSharedMemory(_PosixSharedMemory):
+# FreeBSD (and perhaps other BSDs) limit names to 14 characters.
+_SHM_SAFE_NAME_LENGTH = 14
+
+# shared object name prefix
+if os.name == "nt":
+    _SHM_NAME_PREFIX = 'wnsm_'
+else:
+    _SHM_NAME_PREFIX = '/psm_'
+
+
+def _make_filename():
+    """Create a random filename for the shared memory object.
+    """
+    # number of random bytes to use for name
+    nbytes = (_SHM_SAFE_NAME_LENGTH - len(_SHM_NAME_PREFIX)) // 2
+    assert nbytes >= 2, '_SHM_NAME_PREFIX too long'
+    name = _SHM_NAME_PREFIX + secrets.token_hex(nbytes)
+    assert len(name) <= _SHM_SAFE_NAME_LENGTH
+    return name
+
+
+class PosixSharedMemory:
+
+    # defaults so close() and unlink() can run without errors
+    fd = -1
+    name = None
+    _mmap = None
+    _buf = None
 
     def __init__(self, name, flags=None, mode=384, size=0, read_only=False):
         if name and (flags is None):
-            _PosixSharedMemory.__init__(self, name, mode=mode)
+            flags = 0
         else:
-            if name is None:
-                name = f'psm_{os.getpid()}_{random.randrange(100000)}'
             flags = O_CREX if flags is None else flags
-            _PosixSharedMemory.__init__(
-                self,
-                name,
-                flags=flags,
-                mode=mode,
-                size=size,
-                read_only=read_only
-            )
-
+        if flags & O_EXCL and not flags & O_CREAT:
+            raise ValueError("O_EXCL must be combined with O_CREAT")
+        if name is None and not flags & O_EXCL:
+            raise ValueError("'name' can only be None if O_EXCL is set")
+        flags |= os.O_RDONLY if read_only else os.O_RDWR
+        self.flags = flags
+        self.mode = mode
+        if not size >= 0:
+            raise ValueError("'size' must be a positive integer")
+        if name is None:
+            self._open_retry()
+        else:
+            self.name = name
+            self.fd = _posixshmem.shm_open(name, flags, mode=mode)
+        if size:
+            try:
+                os.ftruncate(self.fd, size)
+            except OSError:
+                self.unlink()
+                raise
         self._mmap = mmap.mmap(self.fd, self.size)
-        self.buf = memoryview(self._mmap)
+        self._buf = memoryview(self._mmap)
+
+    @property
+    def size(self):
+        "Size in bytes."
+        if self.fd >= 0:
+            return os.fstat(self.fd).st_size
+        else:
+            return 0
+
+    @property
+    def buf(self):
+        "A memoryview of contents of the shared memory block."
+        return self._buf
+
+    def _open_retry(self):
+        # generate a random name, open, retry if it exists
+        while True:
+            name = _make_filename()
+            try:
+                self.fd = _posixshmem.shm_open(name, self.flags,
+                                               mode=self.mode)
+            except FileExistsError:
+                continue
+            self.name = name
+            break
 
     def __repr__(self):
         return f'{self.__class__.__name__}({self.name!r}, size={self.size})'
 
+    def unlink(self):
+        if self.name:
+            _posixshmem.shm_unlink(self.name)
+
     def close(self):
-        self.buf.release()
-        self._mmap.close()
-        self.close_fd()
+        if self._buf is not None:
+            self._buf.release()
+            self._buf = None
+        if self._mmap is not None:
+            self._mmap.close()
+            self._mmap = None
+        if self.fd >= 0:
+            os.close(self.fd)
+            self.fd = -1
+
+    def __del__(self):
+        try:
+            self.close()
+        except OSError:
+            pass
 
 
 class SharedMemory:
diff --git a/Lib/test/_test_multiprocessing.py b/Lib/test/_test_multiprocessing.py
@@ -3658,7 +3658,7 @@ def test_shared_memory_basics(self):
             # manages unlinking on its own and unlink() does nothing).
             # True release of shared memory segment does not necessarily
             # happen until process exits, depending on the OS platform.
-            with self.assertRaises(shared_memory.ExistentialError):
+            with self.assertRaises(FileNotFoundError):
                 sms_uno = shared_memory.SharedMemory(
                     'test01_dblunlink',
                     flags=shared_memory.O_CREX,
@@ -3679,7 +3679,7 @@ def test_shared_memory_basics(self):
         # Enforcement of `mode` and `read_only` is OS platform dependent
         # and as such will not be tested here.
 
-        with self.assertRaises(shared_memory.ExistentialError):
+        with self.assertRaises(FileExistsError):
             # Attempting to create a new shared memory segment with a
             # name that is already in use triggers an exception.
             there_can_only_be_one_sms = shared_memory.SharedMemory(
@@ -3704,7 +3704,7 @@ def test_shared_memory_basics(self):
 
         # Attempting to attach to an existing shared memory segment when
         # no segment exists with the supplied name triggers an exception.
-        with self.assertRaises(shared_memory.ExistentialError):
+        with self.assertRaises(FileNotFoundError):
             nonexisting_sms = shared_memory.SharedMemory('test01_notthere')
             nonexisting_sms.unlink()  # Error should occur on prior line.
 
@@ -3745,7 +3745,7 @@ def test_shared_memory_SharedMemoryManager_basics(self):
         if sys.platform != "win32":
             # Calls to unlink() have no effect on Windows platform; shared
             # memory will only be released once final process exits.
-            with self.assertRaises(shared_memory.ExistentialError):
+            with self.assertRaises(FileNotFoundError):
                 # No longer there to be attached to again.
                 absent_shm = shared_memory.SharedMemory(name=held_name)
 
@@ -3756,7 +3756,7 @@ def test_shared_memory_SharedMemoryManager_basics(self):
                 shm = smm2.SharedMemory(size=128)
             held_name = sl.shm.name
         if sys.platform != "win32":
-            with self.assertRaises(shared_memory.ExistentialError):
+            with self.assertRaises(FileNotFoundError):
                 # No longer there to be attached to again.
                 absent_sl = shared_memory.ShareableList(name=held_name)
 
diff --git a/Modules/_multiprocessing/clinic/posixshmem.c.h b/Modules/_multiprocessing/clinic/posixshmem.c.h
@@ -0,0 +1,92 @@
+/*[clinic input]
+preserve
+[clinic start generated code]*/
+
+#if defined(HAVE_SHM_OPEN)
+
+PyDoc_STRVAR(_posixshmem_shm_open__doc__,
+"shm_open($module, /, path, flags, mode=511)\n"
+"--\n"
+"\n"
+"Open a shared memory object.  Returns a file descriptor (integer).");
+
+#define _POSIXSHMEM_SHM_OPEN_METHODDEF    \
+    {"shm_open", (PyCFunction)(void(*)(void))_posixshmem_shm_open, METH_FASTCALL|METH_KEYWORDS, _posixshmem_shm_open__doc__},
+
+static int
+_posixshmem_shm_open_impl(PyObject *module, PyObject *path, int flags,
+                          int mode);
+
+static PyObject *
+_posixshmem_shm_open(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"path", "flags", "mode", NULL};
+    static _PyArg_Parser _parser = {"Ui|i:shm_open", _keywords, 0};
+    PyObject *path;
+    int flags;
+    int mode = 511;
+    int _return_value;
+
+    if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
+        &path, &flags, &mode)) {
+        goto exit;
+    }
+    _return_value = _posixshmem_shm_open_impl(module, path, flags, mode);
+    if ((_return_value == -1) && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = PyLong_FromLong((long)_return_value);
+
+exit:
+    return return_value;
+}
+
+#endif /* defined(HAVE_SHM_OPEN) */
+
+#if defined(HAVE_SHM_UNLINK)
+
+PyDoc_STRVAR(_posixshmem_shm_unlink__doc__,
+"shm_unlink($module, /, path)\n"
+"--\n"
+"\n"
+"Remove a shared memory object (similar to unlink()).\n"
+"\n"
+"Remove a shared memory object name, and, once all processes  have  unmapped\n"
+"the object, de-allocates and destroys the contents of the associated memory\n"
+"region.");
+
+#define _POSIXSHMEM_SHM_UNLINK_METHODDEF    \
+    {"shm_unlink", (PyCFunction)(void(*)(void))_posixshmem_shm_unlink, METH_FASTCALL|METH_KEYWORDS, _posixshmem_shm_unlink__doc__},
+
+static PyObject *
+_posixshmem_shm_unlink_impl(PyObject *module, PyObject *path);
+
+static PyObject *
+_posixshmem_shm_unlink(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"path", NULL};
+    static _PyArg_Parser _parser = {"U:shm_unlink", _keywords, 0};
+    PyObject *path;
+
+    if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
+        &path)) {
+        goto exit;
+    }
+    return_value = _posixshmem_shm_unlink_impl(module, path);
+
+exit:
+    return return_value;
+}
+
+#endif /* defined(HAVE_SHM_UNLINK) */
+
+#ifndef _POSIXSHMEM_SHM_OPEN_METHODDEF
+    #define _POSIXSHMEM_SHM_OPEN_METHODDEF
+#endif /* !defined(_POSIXSHMEM_SHM_OPEN_METHODDEF) */
+
+#ifndef _POSIXSHMEM_SHM_UNLINK_METHODDEF
+    #define _POSIXSHMEM_SHM_UNLINK_METHODDEF
+#endif /* !defined(_POSIXSHMEM_SHM_UNLINK_METHODDEF) */
+/*[clinic end generated code: output=ff9cf0bc9b8baddf input=a9049054013a1b77]*/
diff --git a/Modules/_multiprocessing/posixshmem.c b/Modules/_multiprocessing/posixshmem.c
