Lib/test/test_free_threading/test_io.py

import codecs
import io
import _pyio as pyio
import threading
from unittest import TestCase
from test.support import threading_helper
from test.support.threading_helper import run_concurrently
from random import randint
from sys import getsizeof

threading_helper.requires_working_threading(module=True)


class ThreadSafetyMixin:
    # Test pretty much everything that can break under free-threading.
    # Non-deterministic, but at least one of these things will fail if
    # BytesIO object is not free-thread safe.

    def check(self, funcs, *args):
        barrier = threading.Barrier(len(funcs))
        threads = []

        for func in funcs:
            thread = threading.Thread(target=func, args=(barrier, *args))

            threads.append(thread)

        with threading_helper.start_threads(threads):
            pass

    @threading_helper.requires_working_threading()
    @threading_helper.reap_threads
    def test_free_threading(self):
        """Test for segfaults and aborts."""

        def write(barrier, b, *ignore):
            barrier.wait()
            try: b.write(b'0' * randint(100, 1000))
            except ValueError: pass  # ignore write fail to closed file

        def writelines(barrier, b, *ignore):
            barrier.wait()
            b.write(b'0\n' * randint(100, 1000))

        def truncate(barrier, b, *ignore):
            barrier.wait()
            try: b.truncate(0)
            except BufferError: pass  # ignore exported buffer

        def read(barrier, b, *ignore):
            barrier.wait()
            b.read()

        def read1(barrier, b, *ignore):
            barrier.wait()
            b.read1()

        def readline(barrier, b, *ignore):
            barrier.wait()
            b.readline()

        def readlines(barrier, b, *ignore):
            barrier.wait()
            b.readlines()

        def readinto(barrier, b, into, *ignore):
            barrier.wait()
            b.readinto(into)

        def peek(barrier, b, *ignore):
            barrier.wait()
            b.peek()

        def close(barrier, b, *ignore):
            barrier.wait()
            b.close()

        def getvalue(barrier, b, *ignore):
            barrier.wait()
            b.getvalue()

        def getbuffer(barrier, b, *ignore):
            barrier.wait()
            b.getbuffer()

        def iter(barrier, b, *ignore):
            barrier.wait()
            list(b)

        def getstate(barrier, b, *ignore):
            barrier.wait()
            b.__getstate__()

        def setstate(barrier, b, st, *ignore):
            barrier.wait()
            b.__setstate__(st)

        def sizeof(barrier, b, *ignore):
            barrier.wait()
            getsizeof(b)

        self.check([write] * 10, self.ioclass())
        self.check([writelines] * 10, self.ioclass())
        self.check([write] * 10 + [truncate] * 10, self.ioclass())
        self.check([truncate] + [read] * 10, self.ioclass(b'0\n'*204800))
        self.check([truncate] + [read1] * 10, self.ioclass(b'0\n'*204800))
        self.check([truncate] + [readline] * 10, self.ioclass(b'0\n'*20480))
        self.check([truncate] + [readlines] * 10, self.ioclass(b'0\n'*20480))
        self.check([truncate] + [readinto] * 10, self.ioclass(b'0\n'*204800), bytearray(b'0\n'*204800))
        self.check([truncate] + [peek] * 10, self.ioclass(b'0\n'*204800))
        self.check([close] + [write] * 10, self.ioclass())
        self.check([truncate] + [getvalue] * 10, self.ioclass(b'0\n'*204800))
        self.check([truncate] + [getbuffer] * 10, self.ioclass(b'0\n'*204800))
        self.check([truncate] + [iter] * 10, self.ioclass(b'0\n'*20480))
        self.check([truncate] + [getstate] * 10, self.ioclass(b'0\n'*204800))
        state = self.ioclass(b'123').__getstate__()
        self.check([truncate] + [setstate] * 10, self.ioclass(b'0\n'*204800), state)
        self.check([truncate] + [sizeof] * 10, self.ioclass(b'0\n'*204800))

        # no tests for seek or tell because they don't break anything

class CBytesIOTest(ThreadSafetyMixin, TestCase):
    ioclass = io.BytesIO

    @threading_helper.requires_working_threading()
    @threading_helper.reap_threads
    def test_concurrent_setstate_and_method_call(self):
        # gh-153290: __setstate__() installed the instance __dict__ with a
        # plain store, racing the lock-free LOAD_ATTR method fast path that
        # reads the dict slot with an atomic acquire load.
        states = [(b"A" * 64, 0, {}), (b"B" * 128, 32, {}), (b"C" * 256, 0, {})]
        nreaders = 4
        for _ in range(25):
            shared = self.ioclass(b"initial payload")
            barrier = threading.Barrier(1 + nreaders)

            def setter():
                barrier.wait()
                for state in states:
                    shared.__setstate__(state)

            def reader():
                barrier.wait()
                for _ in range(100):
                    shared.read(8)

            threads = [threading.Thread(target=setter)]
            threads += [threading.Thread(target=reader)
                        for _ in range(nreaders)]
            with threading_helper.start_threads(threads):
                pass

    @threading_helper.requires_working_threading()
    @threading_helper.reap_threads
    def test_concurrent_whole_buffer_read_and_resize(self):
        shared = self.ioclass(b"x" * 64)
        writers = 2
        readers = 8
        loops = 2000
        barrier = threading.Barrier(writers + readers)

        def writer():
            barrier.wait()
            for i in range(loops):
                shared.seek(0)
                shared.write(b"a" * (64 + (i & 63)))

        def reader():
            barrier.wait()
            for _ in range(loops):
                shared.seek(0)
                shared.read()
                shared.seek(0)
                shared.peek()
                shared.getvalue()

        threads = [threading.Thread(target=writer) for _ in range(writers)]
        threads += [threading.Thread(target=reader) for _ in range(readers)]
        with threading_helper.start_threads(threads):
            pass

class PyBytesIOTest(ThreadSafetyMixin, TestCase):
     ioclass = pyio.BytesIO


class IncrementalNewlineDecoderTest(TestCase):
    def make_decoder(self):
        utf8_decoder = codecs.getincrementaldecoder('utf-8')()
        return io.IncrementalNewlineDecoder(utf8_decoder, translate=True)

    def test_concurrent_reset(self):
        decoder = self.make_decoder()

        def worker():
            for _ in range(100):
                decoder.reset()

        run_concurrently(worker_func=worker, nthreads=2)

    def test_concurrent_decode(self):
        decoder = self.make_decoder()

        def worker():
            for _ in range(100):
                decoder.decode(b"line\r\n", final=False)

        run_concurrently(worker_func=worker, nthreads=2)

    def test_concurrent_getstate_setstate(self):
        decoder = self.make_decoder()
        state = decoder.getstate()

        def getstate_worker():
            for _ in range(100):
                decoder.getstate()

        def setstate_worker():
            for _ in range(100):
                decoder.setstate(state)

        run_concurrently([getstate_worker] * 2 + [setstate_worker] * 2)

    def test_concurrent_decode_and_reset(self):
        decoder = self.make_decoder()

        def decode_worker():
            for _ in range(100):
                decoder.decode(b"line\r\n", final=False)

        def reset_worker():
            for _ in range(100):
                decoder.reset()

        run_concurrently([decode_worker] * 2 + [reset_worker] * 2)


class TextIOWrapperTest(TestCase):
    def test_buffer_detach_race(self):
        make = lambda: io.TextIOWrapper(io.BytesIO())
        slot = [make()]

        def reader():
            for _ in range(1000):
                try:
                    slot[0].buffer
                except ValueError:
                    pass

        def detacher():
            for _ in range(1000):
                slot[0] = make()
                slot[0].detach()

        run_concurrently([reader, detacher])