Decide threads and the GIL per chunk for Python codecs and filters

src/blosc2/blosc2_ext.pyx

@@ -199,6 +199,13 @@ cdef extern from "blosc2.h":
         BLOSC2_SPECIAL_VALUE
         BLOSC2_SPECIAL_UNINIT
 
+    ctypedef enum:
+        BLOSC_DOSHUFFLE
+        BLOSC_DOBITSHUFFLE
+        BLOSC_UDCODEC_FORMAT
+        BLOSC2_CHUNK_FLAGS
+        BLOSC2_CHUNK_FILTER_CODES
+
     ctypedef enum:
         BLOSC2_VERSION_STRING
         BLOSC2_VERSION_REVISION
@@ -3374,18 +3381,19 @@ def decompress(src, dst=None, as_bytearray=False):
     typed_view_src = mem_view_src.cast('B')
     _check_comp_length('src', len(typed_view_src))
     blosc2_cbuffer_sizes(<void*>&typed_view_src[0], &nbytes, &cbytes, &blocksize)
+    cdef int32_t srcsize = <int32_t>len(typed_view_src)
     cdef Py_buffer buf
     if dst is not None:
         PyObject_GetBuffer(dst, &buf, PyBUF_SIMPLE)
         if buf.len == 0:
             raise ValueError("The dst length must be greater than 0")
-        size = blosc1_decompress(<void*>&typed_view_src[0], buf.buf, buf.len)
+        size = _blosc1_decompress(&typed_view_src[0], srcsize, buf.buf, buf.len)
         PyBuffer_Release(&buf)
     else:
         dst = PyBytes_FromStringAndSize(NULL, nbytes)
         if dst is None:
             raise RuntimeError("Could not get a bytes object")
-        size = blosc1_decompress(<void*>&typed_view_src[0], <void*> <char *> dst, len(dst))
+        size = _blosc1_decompress(&typed_view_src[0], srcsize, <void*> <char *> dst, len(dst))
         if size >= 0:
             return bytearray(dst) if as_bytearray else dst
     if size < 0:
@@ -3468,8 +3476,9 @@ def set_releasegil(c_bool gilstate):
 
 # Ids of the codecs and filters registered from Python (register_codec(), register_filter()), whose
 # callbacks need the GIL.  C-Blosc2 cannot unregister them, so these are never cleared.  They let the GIL
-# checks of compressions run without Python lookups; slice reads (_can_release_gil()) still look at
-# blosc2.ucodecs_registry and blosc2.ufilters_registry, which tests replace to force reads without the GIL.
+# checks of compressions and chunk headers run without Python lookups; slice reads (_can_release_gil())
+# still look at blosc2.ucodecs_registry and blosc2.ufilters_registry, which tests replace to force reads
+# without the GIL.
 cdef c_bool _python_codec_ids[256]
 cdef c_bool _python_filter_ids[256]
 cdef c_bool _python_ids_registered = False
@@ -3513,9 +3522,9 @@ cdef int _compress_ctx(blosc2_context *cctx, const void *src, int32_t srcsize, v
 cdef int _decompress_ctx(blosc2_context *dctx, const void *src, int32_t srcsize, void *dest,
                          int32_t destsize) noexcept:
     """blosc2_decompress_ctx() with a context without postfilter, without the GIL if set_releasegil() asks
-    and no codec or filter is registered from Python (the chunk tells which ones it uses)."""
+    and the chunk uses no codec or filter registered from Python."""
     cdef int rc
-    if RELEASEGIL and not _python_ids_registered:
+    if RELEASEGIL and not _chunk_runs_python(<const uint8_t*>src, srcsize):
         with nogil:
             rc = blosc2_decompress_ctx(dctx, src, srcsize, dest, destsize)
     else:
@@ -3704,15 +3713,54 @@ cdef create_dparams_from_kwargs(blosc2_dparams *dparams, kwargs, blosc2_cparams*
     #dparams.typesize = typesize = kwargs.get('typesize', blosc2.dparams_dflts['typesize'])
     _check_dparams(dparams, cparams)
 
+# The offset in the extended header of a chunk of the id of its codec when the codec bits of its flags say
+# BLOSC_UDCODEC_FORMAT (README_CHUNK_FORMAT.rst in C-Blosc2; blosc2.h has no name for it)
+cdef enum:
+    _CHUNK_UDCODEC = 0x16
+
+cdef c_bool _chunk_runs_python(const uint8_t *chunk, int32_t chunksize) noexcept:
+    """Whether decompressing `chunk` can run a codec or filter registered from Python, as its header tells.
+    Their callbacks need the GIL, so they must run on the calling thread, holding it: C-Blosc2's worker
+    threads never hold it."""
+    cdef uint8_t flags
+    cdef uint8_t codec
+    cdef int i
+    if not _python_ids_registered or chunksize < BLOSC_EXTENDED_HEADER_LENGTH:
+        return False
+    flags = chunk[BLOSC2_CHUNK_FLAGS]
+    # Only an extended header (both shuffle flags set) can name user codecs and filters
+    if not (flags & BLOSC_DOSHUFFLE and flags & BLOSC_DOBITSHUFFLE):
+        return False
+    codec = flags >> 5  # the codec bits of the flags (5-7)
+    if codec == BLOSC_UDCODEC_FORMAT:
+        codec = chunk[_CHUNK_UDCODEC]
+    if _python_codec_ids[codec]:
+        return True
+    for i in range(BLOSC2_MAX_FILTERS):
+        if _python_filter_ids[chunk[BLOSC2_CHUNK_FILTER_CODES + i]]:
+            return True
+    return False
+
+cdef int _blosc1_decompress(const uint8_t *src, int32_t srcsize, void *dest, size_t destsize) noexcept:
+    """blosc1_decompress(), which uses the global context and its number of threads, except for chunks that
+    can run Python code (see _chunk_runs_python()): those get a context of one thread of their own."""
+    cdef blosc2_context *dctx
+    cdef int rc
+    if not _chunk_runs_python(src, srcsize):
+        return blosc1_decompress(src, dest, destsize)
+    dctx = blosc2_create_dctx(BLOSC2_DPARAMS_DEFAULTS)  # one thread
+    if dctx == NULL:
+        return BLOSC2_ERROR_MEMORY_ALLOC
+    rc = blosc2_decompress_ctx(dctx, src, srcsize, dest, <int32_t>destsize)
+    blosc2_free_ctx(dctx)
+    return rc
+
 def decompress2(src, dst=None, **kwargs):
     cdef blosc2_dparams dparams
     cdef char *dst_buf
     cdef void *view
     create_dparams_from_kwargs(&dparams, kwargs)
 
-    cdef blosc2_context *dctx = blosc2_create_dctx(dparams)
-    if dctx == NULL:
-        raise RuntimeError("Could not create decompression context")
     cdef const uint8_t[:] typed_view_src
     mem_view_src = memoryview(src)
     typed_view_src = mem_view_src.cast('B')
@@ -3722,6 +3770,11 @@ def decompress2(src, dst=None, **kwargs):
     cdef int32_t blocksize
     cdef int32_t srcsize = <int32_t>typed_view_src.nbytes
     blosc2_cbuffer_sizes(<void*>&typed_view_src[0], &nbytes, &cbytes, &blocksize)
+    if dparams.nthreads > 1 and _chunk_runs_python(&typed_view_src[0], srcsize):
+        dparams.nthreads = 1
+    cdef blosc2_context *dctx = blosc2_create_dctx(dparams)
+    if dctx == NULL:
+        raise RuntimeError("Could not create decompression context")
     cdef Py_buffer buf
     if dst is not None:
         PyObject_GetBuffer(dst, &buf, PyBUF_SIMPLE)
@@ -3824,10 +3877,6 @@ def vldecompress(src, **kwargs):
     cdef blosc2_dparams dparams
     create_dparams_from_kwargs(&dparams, kwargs)
 
-    cdef blosc2_context *dctx = blosc2_create_dctx(dparams)
-    if dctx == NULL:
-        raise RuntimeError("Could not create decompression context")
-
     cdef const uint8_t[:] typed_view_src
     mem_view_src = memoryview(src)
     typed_view_src = mem_view_src.cast('B')
@@ -3837,6 +3886,11 @@ def vldecompress(src, **kwargs):
     cdef int32_t nblocks
     cdef int32_t srcsize = <int32_t>typed_view_src.nbytes
     blosc2_cbuffer_sizes(<void*>&typed_view_src[0], &nbytes, &cbytes, &nblocks)
+    if dparams.nthreads > 1 and _chunk_runs_python(&typed_view_src[0], srcsize):
+        dparams.nthreads = 1
+    cdef blosc2_context *dctx = blosc2_create_dctx(dparams)
+    if dctx == NULL:
+        raise RuntimeError("Could not create decompression context")
     if nblocks <= 0:
         blosc2_free_ctx(dctx)
         raise ValueError("Chunk does not contain VL blocks")
@@ -4702,9 +4756,7 @@ cdef class SChunk:
 
     @_handle_locked
     def to_cframe(self):
-        # Serializing can decompress (with the postfilter) and recompress chunks and vlmetalayers
-        return schunk_to_cframe(self.schunk, _can_compress_without_gil(self.schunk.cctx)
-                                and _can_release_gil(self.schunk))
+        return schunk_to_cframe(self.schunk, _can_serialize_without_gil(self.schunk))
 
     @_handle_locked
     def _avoid_cframe_free(self, avoid_cframe_free):
@@ -5973,6 +6025,31 @@ cdef schunk_is_ndarray(blosc2_schunk* schunk):
     return blosc2_meta_exists(schunk, meta) >= 0
 
 
+cdef c_bool _can_serialize_without_gil(blosc2_schunk *sc) except? False:
+    """Whether to_cframe() may release the GIL as set_releasegil() asks: only if serializing `sc` cannot run
+    a Python callback.  In memory (see blosc2_schunk_to_buffer_alloc()), a contiguous frame is only copied,
+    and a sparse super-chunk without prefilter and postfilter gets its chunks copied as they are and its
+    vlmetalayers decompressed and recompressed with the default cparams, so only their codecs and filters
+    matter.  Anything else keeps the check of compressions and slice reads: a prefilter or postfilter makes
+    C-Blosc2 decompress and recompress every chunk, and a frame on disk can get its vlmetalayers reloaded
+    during the call."""
+    cdef blosc2_cparams cparams
+    cdef blosc2_dparams dparams
+    cdef int i
+    if RELEASEGIL and sc.storage.urlpath == NULL:
+        if sc.frame != NULL and sc.storage.contiguous:
+            return True
+        blosc2_ctx_get_cparams(sc.cctx, &cparams)
+        blosc2_ctx_get_dparams(sc.dctx, &dparams)
+        if (sc.frame == NULL and cparams.prefilter == NULL and dparams.postfilter == NULL
+                and sc.storage.dparams.postfilter == NULL):
+            for i in range(sc.nvlmetalayers):
+                if _chunk_runs_python(sc.vlmetalayers[i].content, sc.vlmetalayers[i].content_len):
+                    return False
+            return True
+    return _can_compress_without_gil(sc.cctx) and _can_release_gil(sc)
+
+
 cdef struct cframe_dest:
     PyObject* out
     c_bool alloc_failed

src/blosc2/core.py

@@ -1529,8 +1529,8 @@ def set_releasegil(gilstate: bool) -> bool:
 
     Compressing keeps the GIL regardless of this setting when a Python callback
     (a prefilter, a filler, or a codec or filter registered from Python) could
-    run.  Decompressing a chunk keeps it while any codec or filter is registered
-    from Python, as the chunk may use one.
+    run.  :func:`decompress2` keeps it when the chunk uses a codec or filter
+    registered from Python.
 
     Examples
     --------

tests/test_python_codec_threads.py

@@ -0,0 +1,119 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org>
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#######################################################################
+
+# Codecs and filters registered from Python need the GIL, which C-Blosc2's worker threads never hold, so
+# chunks that use them must decompress on the calling thread, whatever number of threads is asked for.
+# Running them on a worker thread crashes the interpreter: the decompressions run in a subprocess.
+
+import json
+import os
+import subprocess
+import sys
+import zlib
+
+import numpy as np
+import pytest
+
+import blosc2
+
+CODEC_ID = 170
+FILTER_ID = 171
+N = 1000  # int32 items per chunk, in 4 blocks, so that several threads can share a chunk
+# Offsets in a chunk header: flags (BLOSC2_CHUNK_FLAGS), filter ids (BLOSC2_CHUNK_FILTER_CODES) and the
+# id of a user-defined codec, which the codec bits of the flags (the top 3) name with BLOSC_UDCODEC_FORMAT
+CHUNK_FLAGS = 0x2
+CHUNK_FILTER_CODES = 0x10
+CHUNK_UDCODEC = 0x16
+UDCODEC_FORMAT = 6
+
+
+def encoder(input, output, meta, schunk):
+    out = zlib.compress(input.tobytes())
+    if len(out) >= output.size:
+        return 0
+    output[: len(out)] = np.frombuffer(out, dtype=np.uint8)
+    return len(out)
+
+
+def decoder(input, output, meta, schunk):
+    out = zlib.decompress(input.tobytes())
+    output[: len(out)] = np.frombuffer(out, dtype=np.uint8)
+    return len(out)
+
+
+def xor(input, output, meta, schunk):
+    np.bitwise_xor(input, 0x5A, out=output)
+
+
+def outcome(func, *args, **kwargs):
+    try:
+        func(*args, **kwargs)
+    except Exception as e:
+        return f"{type(e).__name__}: {e} (from {e.__cause__ or e.__context__!r})"
+    return "ok"
+
+
+def decompression_outcomes(nthreads):
+    """How decompressing chunks that use the Python codec or filter outside of a super-chunk ends, with
+    `nthreads` threads; and whether a chunk that uses neither still decompresses."""
+    blosc2.register_codec("test_threads", CODEC_ID, encoder, decoder)
+    blosc2.register_filter(FILTER_ID, xor, xor, "test_threads")
+    blosc2.set_nthreads(nthreads)  # decompress() uses the global number of threads
+    values = (np.arange(N) % 200).astype(np.int32)
+    cparams = {"typesize": values.itemsize, "nthreads": 1, "blocksize": values.nbytes // 4}
+    dparams = {"nthreads": 1}  # super-chunks refuse more with Python codecs and filters
+    python_cparams = {
+        "codec": {**cparams, "codec": CODEC_ID, "filters": [blosc2.Filter.NOFILTER], "filters_meta": [0]},
+        "filter": {**cparams, "filters": [FILTER_ID], "filters_meta": [0]},
+    }
+    results = {}
+    for kind, kind_cparams in python_cparams.items():
+        schunk = blosc2.SChunk(chunksize=values.nbytes, data=values, cparams=kind_cparams, dparams=dparams)
+        chunk = schunk.get_chunk(0)
+        results[f"decompress2 {kind}"] = outcome(blosc2.decompress2, chunk, nthreads=nthreads)
+        results[f"decompress {kind}"] = outcome(blosc2.decompress, chunk)
+    # vlcompress() cannot run the Python codec or filter (no super-chunk): name them in the header of a
+    # plain VL-block chunk instead
+    vlchunk = blosc2.blosc2_ext.vlcompress([values.tobytes()] * 4, **cparams)
+    vlchunks = {"codec": bytearray(vlchunk), "filter": bytearray(vlchunk)}
+    vlchunks["codec"][CHUNK_FLAGS] = vlchunk[CHUNK_FLAGS] & 0x1F | UDCODEC_FORMAT << 5
+    vlchunks["codec"][CHUNK_UDCODEC] = CODEC_ID
+    vlchunks["filter"][CHUNK_FILTER_CODES] = FILTER_ID
+    for kind, chunk in vlchunks.items():
+        results[f"vldecompress {kind}"] = outcome(
+            blosc2.blosc2_ext.vldecompress, bytes(chunk), nthreads=nthreads
+        )
+    plain = blosc2.compress2(values, **cparams)
+    results["decompress2 plain"] = blosc2.decompress2(plain, nthreads=nthreads) == values.tobytes()
+    return results
+
+
+@pytest.mark.skipif(blosc2.IS_WASM, reason="no subprocesses and no threads on wasm32")
+def test_python_codec_filter_decompress_threads():
+    env = {name: value for name, value in os.environ.items() if name != "BLOSC_NTHREADS"}
+    runs = {}
+    for nthreads in (1, 4):
+        proc = subprocess.run(
+            [sys.executable, __file__, str(nthreads)],
+            capture_output=True,
+            text=True,
+            env=env,
+            check=False,
+            timeout=300,
+        )
+        assert proc.returncode == 0, f"{nthreads} threads: exit code {proc.returncode}\n{proc.stderr}"
+        runs[nthreads] = json.loads(proc.stdout.splitlines()[-1])
+    # The callbacks raise this themselves: Python codecs and filters only work in a super-chunk
+    assert all("without an SChunk" in result for call, result in runs[1].items() if "plain" not in call), (
+        runs[1]
+    )
+    assert runs[1]["decompress2 plain"] is True, runs[1]
+    assert runs[4] == runs[1]
+
+
+if __name__ == "__main__":
+    print(json.dumps(decompression_outcomes(int(sys.argv[1]))))