gh-155742: Use PyBytesWriter in marshal - #155748

I do not think the current code is broken.

I didn't say that the current code is broken. The PR only just avoids the soft deprecated PyBytes_FromStringAndSize() function.

I reworked the error handling.

@serhiy-storchaka: Please review the updated PR.

I wrote a script to test manually this PR by injecting MemoryError at different places:

import marshal
import io
import _testcapi
obj = b'x' * (1024 * 1024)
file = io.BytesIO()
for i in range(10):
    try:
        try:
            _testcapi.set_nomemory(i)
            res = marshal.dump(obj, file)
        finally:
            _testcapi.remove_mem_hooks()
    except Exception as exc:
        print(f"marshal.dump failed: {exc!r}")
    else:
        print(f"{res=}")

Before, the code failed with an assertion error. With my latest change, the code works is all cases (always raise MemoryError as expected).

Ah, I noticed that the PyMarshal C API is not tested by test_capi currently. So I wrote PR gh-156890 to add tests.

I extracted the TYPE_STRING change: it does in fact fix an issue, using PyBytes_FromStringAndSize(str, n) allows getting 1-byte singletons: PR gh-157398.

I do not think we need this change. It looks to me like a code churn which makes the code more complicated.

UPDATE: Oh sorry, at my first attempt, I ran benchmarks on a debug build! I replaced results with a benchmark on a release build.

I ran a quick benchmark on marshal.dumps():

import pyperf
import marshal

def noop_func():
    pass

runner = pyperf.Runner()
for obj in (b'abc', True, 123):
    runner.bench_func(f'dumps {obj!r}', marshal.dumps, obj)

runner.bench_func('dumps code object', marshal.dumps, noop_func.__code__)
runner.bench_func('dumps list(range(20))', marshal.dumps, list(range(20)))
runner.bench_func("dumps '\u20ac' * 10", marshal.dumps, '\u20ac' * 10)
runner.bench_func("dumps 'long line '*1000", marshal.dumps, 'long line '*1000)

Results:

I ran a second benchmark building code objects of the stdlib top 10 largest files:

import pyperf
import marshal
import tokenize

# Top 10 largest stdlib Python files
files = (
    'subprocess.py',
    'pydoc.py',
    'doctest.py',
    'argparse.py',
    'tarfile.py',
    'inspect.py',
    'typing.py',
    'pdb.py',
    'turtle.py',
    '_pydecimal.py',
)

runner = pyperf.Runner()
for filename in files:
    with tokenize.open("Lib/" + filename) as fp:
        code = fp.read()
        code = compile(code, filename, "exec")
        runner.bench_func(f'dumps {filename}', marshal.dumps, code)

Results:

Benchmark hidden because not significant (2): dumps tarfile.py, dumps _pydecimal.py

Aha, so using marshal.c overallocation (delta = size + 1024) + PyBytesWriter overallocation is less efficient than the current code. I made a tiny change: disable PyBytesWriter overallocation, and now this change makes marshal.dumps() faster.

First benchmark on simple small objects:

Second benchmark on large objects from stdlib modules:

I'm surprised that this change makes marshal.dumps() faster. I expected same performance or slower. So I reran the benchmark with CPU isolation. It's still faster on all benchmarks, except of 'long line '*1000 (1.01x slower, minor difference in fact).

First benchmark on simple small objects:

Second benchmark on large objects from stdlib modules:

These results are on the latest version of this PR, where the latest commit disables PyBytesWriter overallocation.

marshal.dumps(True) is a single byte (b'T').

• Currently, dumps() starts with PyBytes_FromStringAndSize(50) and then calls _PyBytes_Resize(1).
• Using PyBytesWriter, it uses the PyBytesWriter small buffer (256 bytes) and then calls PyBytes_FromStringAndSize(1) which doesn't even allocate an object but returns a singleton. That explains why this case is the fastest (1.21x faster).
• Previously, the bytes object had to be resized from 50 to 1 bytes. With this change, no bytes object is created.
• In the common case, PyBytesWriter_Create() doesn't allocate heap memory for the writer, but use the free list, so the allocation is very fast.

I compared the PyBytes calls before/after on _pydecimal.py (1.10x faster)

Before:

• marshal: 0 => PyBytes_FromStringAndSize(50)
• w_reserve(): 50 => _PyBytes_Resize(2096)
• w_reserve(): 2096 => _PyBytes_Resize(5216)
• w_reserve(): 5216 => _PyBytes_Resize(11456)
• w_reserve(): 11456 => _PyBytes_Resize(23936)
• w_reserve(): 23936 => _PyBytes_Resize(48896)
• w_reserve(): 48896 => _PyBytes_Resize(98816)
• w_reserve(): 98816 => _PyBytes_Resize(198656)
• w_reserve(): 198656 => _PyBytes_Resize(398336)
• marshal final: 398336 => _PyBytes_Resize(228842)
• dumps(): 8 calls to w_reserve()

After:

• --- use PyBytesWriter small buffer (256 bytes) ---
• PyBytesWriter_Resize(): 255 => PyBytes_FromStringAndSize(2096)
• PyBytesWriter_Resize(): 2096 => _PyBytes_Resize(5216)
• PyBytesWriter_Resize(): 5216 => _PyBytes_Resize(11456)
• PyBytesWriter_Resize(): 11456 => _PyBytes_Resize(23936)
• PyBytesWriter_Resize(): 23936 => _PyBytes_Resize(48896)
• PyBytesWriter_Resize(): 48896 => _PyBytes_Resize(98816)
• PyBytesWriter_Resize(): 98816 => _PyBytes_Resize(198656)
• PyBytesWriter_Resize(): 198656 => _PyBytes_Resize(398336)
• PyBytesWriter_FinishWithSize(): 398336 => _PyBytes_Resize(228842)
• dumps(): 8 calls to w_reserve()

The only difference is that using PyBytesWriter, the bytes object is created with 2096 bytes, whereas it's created with 50 bytes currently.

I checked the PyBytes calls before/after on subprocess.py (1.08x faster).

To serialize {"emscripten", "wasi", "ios", "tvos", "watchos"} literal set, w_complex_object() calls _PyMarshal_WriteObjectToString() on each string. Since PyBytesWriter is faster for small objects (up to 256 bytes), it serializes the set literal faster.

Before:

• marshal: 0 => PyBytes_FromStringAndSize(50)
• w_reserve(): 50 => _PyBytes_Resize(1128)
• w_reserve(): 1128 => _PyBytes_Resize(3280)
• w_reserve(): 3280 => _PyBytes_Resize(7584)
• w_reserve(): 7584 => _PyBytes_Resize(16192)
• w_reserve(): 16192 => _PyBytes_Resize(33408)
• w_reserve(): 33408 => _PyBytes_Resize(67840)
• w_reserve(): 67840 => _PyBytes_Resize(136704)
  • marshal: 0 => PyBytes_FromStringAndSize(50) -- nested call to serialize a set member (str)
  • marshal final: 50 => _PyBytes_Resize(12)
  • marshal: 0 => PyBytes_FromStringAndSize(50) -- nested call to serialize a set member (str)
  • marshal final: 50 => _PyBytes_Resize(9)
  • marshal: 0 => PyBytes_FromStringAndSize(50) -- nested call to serialize a set member (str)
  • marshal final: 50 => _PyBytes_Resize(6)
  • marshal: 0 => PyBytes_FromStringAndSize(50) -- nested call to serialize a set member (str)
  • marshal final: 50 => _PyBytes_Resize(5)
  • marshal: 0 => PyBytes_FromStringAndSize(50) -- nested call to serialize a set member (str)
  • marshal final: 50 => _PyBytes_Resize(6)
• marshal final: 136704 => _PyBytes_Resize(94617)
• dumps(): 7 calls to w_reserve()

After:

• PyBytesWriter_Resize(): 255 => PyBytes_FromStringAndSize(1128)
• PyBytesWriter_Resize(): 1128 => _PyBytes_Resize(3280)
• PyBytesWriter_Resize(): 3280 => _PyBytes_Resize(7584)
• PyBytesWriter_Resize(): 7584 => _PyBytes_Resize(16192)
• PyBytesWriter_Resize(): 16192 => _PyBytes_Resize(33408)
• PyBytesWriter_Resize(): 33408 => _PyBytes_Resize(67840)
• PyBytesWriter_Resize(): 67840 => _PyBytes_Resize(136704)
• PyBytesWriter_FinishWithSize(): 136704 => _PyBytes_Resize(94617)
• dumps(): 7 calls to w_reserve()

I tried to reproduce the benchmark: release build of main vs main + this PR (07d4266), pinned to one core, min of 15 repeats, 3 interleaved rounds.

The only measurable win is for tiny outputs like True (67 → 60 ns), where the 50-byte bytes object is no longer allocated and shrunk. For code objects of stdlib modules the difference is within noise (if anything, 1–2% slower). This agrees with your own trace: the sequence of _PyBytes_Resize() calls is identical except for the first allocation, so there is no reason to expect a 10% difference.

So the performance argument does not hold, and the code is not simpler. wf.writer->overallocate = 0 writes to a private field of a structure which is opaque in the public API. That marshal has to disable the writer's growth policy to keep its own shows that PyBytesWriter is a poor fit here: it is used only as a holder for a bytes object plus a 256-byte stack buffer.

Other notes, in case you want to pursue this anyway:

• The final PyBytesWriter_Resize(wf.writer, wf.ptr - base) does not shrink the buffer, PyBytesWriter_Finish() does. PyBytesWriter_FinishWithPointer(wf.writer, wf.ptr) would do the same in one call, and there would be no need to resize before discarding on the error path.
• PyBytesWriter_Create(50) uses only 50 of the 256 bytes of the small buffer. Any output between 51 and 256 bytes still allocates a 1124-byte bytes object and then shrinks it.

I still think that this change is not needed.