test_resize_race.py

#!/usr/bin/env python3
"""
Targeted tests to narrow down the resize race condition.

The lock-free implementation does this during resize:
1. FT_ATOMIC_STORE_PTR_RELEASE(so->table, NULL)
2. FT_ATOMIC_STORE_SSIZE_RELAXED(so->mask, newsize - 1)
3. ... copy entries ...
4. FT_ATOMIC_STORE_PTR_RELEASE(so->table, newtable)

And set_lookkey_threadsafe does:
1. table = FT_ATOMIC_LOAD_PTR_ACQUIRE(so->table)
2. mask = FT_ATOMIC_LOAD_SSIZE_RELAXED(so->mask)
3. if (table == NULL || table != so->table) -> fallback

Potential issues:
- mask could be read after table, getting new mask with old table
- QSBR cleanup might free table while reader is using it
- set_empty_to_minsize also sets table=NULL temporarily
"""

import argparse
import sys
import threading
import time
from dataclasses import dataclass


@dataclass
class Stats:
    ops: int = 0
    errors: int = 0
    false_negatives: int = 0  # Key should be present but wasn't found
    false_positives: int = 0  # Key shouldn't be present but was found


def test_resize_during_contains():
    """
    Test 1: Single writer grows set, readers check contains.
    No clear() - just growth causing resizes.
    """
    print("\n=== Test 1: Resize during contains (grow only) ===")

    s = set()
    stop = threading.Event()
    barrier = threading.Barrier(5)
    reader_stats = [Stats() for _ in range(4)]
    writer_stats = Stats()

    # Track what values have been added (for verification)
    added_values = set()
    added_lock = threading.Lock()

    def reader(tid):
        stats = reader_stats[tid]
        barrier.wait()
        while not stop.is_set():
            # Check for values we know have been added
            with added_lock:
                check_values = list(added_values)[:100]  # Sample

            for v in check_values:
                result = v in s
                stats.ops += 1
                # Value was added, should be found (unless cleared)
                if not result:
                    stats.false_negatives += 1

    def writer():
        nonlocal added_values
        barrier.wait()
        val = 0
        while not stop.is_set():
            s.add(val)
            with added_lock:
                added_values.add(val)
            writer_stats.ops += 1
            val += 1
            # Don't clear - just grow to trigger resizes

    threads = [threading.Thread(target=reader, args=(i,)) for i in range(4)]
    threads.append(threading.Thread(target=writer))

    for t in threads:
        t.start()

    time.sleep(5)
    stop.set()

    for t in threads:
        t.join(timeout=2)

    total_reader_ops = sum(st.ops for st in reader_stats)
    total_false_neg = sum(st.false_negatives for st in reader_stats)

    print(f"  Writer ops: {writer_stats.ops:,}")
    print(f"  Reader ops: {total_reader_ops:,}")
    print(f"  False negatives: {total_false_neg}")
    print(f"  Final set size: {len(s)}")

    # Check final state
    missing = added_values - s
    if missing:
        print(f"  ERROR: {len(missing)} values missing from final set!")
        return False

    return total_false_neg == 0


def test_clear_during_contains():
    """
    Test 2: Writer clears set, readers check contains.
    Tests set_empty_to_minsize which also sets table=NULL.
    """
    print("\n=== Test 2: Clear during contains ===")

    s = set()
    stop = threading.Event()
    barrier = threading.Barrier(5)
    reader_stats = [Stats() for _ in range(4)]
    writer_stats = Stats()

    def reader(tid):
        stats = reader_stats[tid]
        barrier.wait()
        while not stop.is_set():
            for i in range(100):
                _ = i in s
                stats.ops += 1

    def writer():
        barrier.wait()
        while not stop.is_set():
            # Add values then clear
            for i in range(50):
                s.add(i)
                writer_stats.ops += 1
            s.clear()
            writer_stats.ops += 1

    threads = [threading.Thread(target=reader, args=(i,)) for i in range(4)]
    threads.append(threading.Thread(target=writer))

    for t in threads:
        t.start()

    time.sleep(5)
    stop.set()

    for t in threads:
        t.join(timeout=2)

    total_reader_ops = sum(st.ops for st in reader_stats)
    total_errors = sum(st.errors for st in reader_stats)

    print(f"  Writer ops: {writer_stats.ops:,}")
    print(f"  Reader ops: {total_reader_ops:,}")
    print(f"  Errors: {total_errors}")

    return total_errors == 0


def test_threshold_crossing():
    """
    Test 3: Repeatedly cross the PySet_MINSIZE (8) threshold.
    This switches between smalltable and malloced table.
    """
    print("\n=== Test 3: Threshold crossing (smalltable <-> malloc) ===")

    s = set()
    stop = threading.Event()
    barrier = threading.Barrier(5)
    reader_stats = [Stats() for _ in range(4)]
    writer_stats = Stats()
    errors = []
    error_lock = threading.Lock()

    def reader(tid):
        stats = reader_stats[tid]
        barrier.wait()
        while not stop.is_set():
            try:
                for i in range(20):
                    _ = i in s
                    stats.ops += 1
            except Exception as e:
                with error_lock:
                    errors.append(f"Reader {tid}: {e}")
                stats.errors += 1

    def writer():
        barrier.wait()
        while not stop.is_set():
            try:
                # Grow past threshold (8)
                for i in range(12):
                    s.add(i)
                    writer_stats.ops += 1

                # Shrink below threshold using pop
                while len(s) > 4:
                    s.pop()
                    writer_stats.ops += 1
            except Exception as e:
                with error_lock:
                    errors.append(f"Writer: {e}")

    threads = [threading.Thread(target=reader, args=(i,)) for i in range(4)]
    threads.append(threading.Thread(target=writer))

    for t in threads:
        t.start()

    time.sleep(5)
    stop.set()

    for t in threads:
        t.join(timeout=2)

    total_reader_ops = sum(st.ops for st in reader_stats)

    print(f"  Writer ops: {writer_stats.ops:,}")
    print(f"  Reader ops: {total_reader_ops:,}")
    print(f"  Errors: {len(errors)}")
    if errors:
        for e in errors[:5]:
            print(f"    {e}")

    return len(errors) == 0


def test_multiple_writers_clear():
    """
    Test 4: Multiple writers all clearing - high contention on resize.
    """
    print("\n=== Test 4: Multiple writers with clear ===")

    s = set()
    stop = threading.Event()
    barrier = threading.Barrier(6)  # 4 readers + 2 writers
    reader_stats = [Stats() for _ in range(4)]
    writer_stats = [Stats() for _ in range(2)]
    errors = []
    error_lock = threading.Lock()

    def reader(tid):
        stats = reader_stats[tid]
        barrier.wait()
        while not stop.is_set():
            try:
                for i in range(50):
                    _ = i in s
                    stats.ops += 1
            except Exception as e:
                with error_lock:
                    errors.append(f"Reader {tid}: {e}")
                stats.errors += 1

    def writer(tid):
        stats = writer_stats[tid]
        barrier.wait()
        base = tid * 100
        while not stop.is_set():
            try:
                for i in range(20):
                    s.add(base + i)
                    stats.ops += 1
                s.clear()
                stats.ops += 1
            except Exception as e:
                with error_lock:
                    errors.append(f"Writer {tid}: {e}")

    threads = [threading.Thread(target=reader, args=(i,)) for i in range(4)]
    threads.extend(
        [threading.Thread(target=writer, args=(i,)) for i in range(2)]
    )

    for t in threads:
        t.start()

    time.sleep(5)
    stop.set()

    for t in threads:
        t.join(timeout=2)

    total_reader_ops = sum(st.ops for st in reader_stats)
    total_writer_ops = sum(st.ops for st in writer_stats)

    print(f"  Writer ops: {total_writer_ops:,}")
    print(f"  Reader ops: {total_reader_ops:,}")
    print(f"  Errors: {len(errors)}")
    if errors:
        for e in errors[:5]:
            print(f"    {e}")

    return len(errors) == 0


def test_rapid_resize():
    """
    Test 5: Maximize resize frequency by using update() and clear().
    """
    print("\n=== Test 5: Rapid resize with update/clear ===")

    s = set()
    stop = threading.Event()
    barrier = threading.Barrier(5)
    reader_stats = [Stats() for _ in range(4)]
    writer_stats = Stats()
    errors = []
    error_lock = threading.Lock()

    def reader(tid):
        stats = reader_stats[tid]
        barrier.wait()
        while not stop.is_set():
            try:
                # Tight loop of contains
                for i in range(100):
                    _ = i in s
                    stats.ops += 1
            except Exception as e:
                with error_lock:
                    errors.append(f"Reader {tid}: {e}")
                stats.errors += 1

    def writer():
        barrier.wait()
        while not stop.is_set():
            try:
                # update() causes resize if needed
                s.update(range(100))
                writer_stats.ops += 1
                s.clear()
                writer_stats.ops += 1
            except Exception as e:
                with error_lock:
                    errors.append(f"Writer: {e}")

    threads = [threading.Thread(target=reader, args=(i,)) for i in range(4)]
    threads.append(threading.Thread(target=writer))

    for t in threads:
        t.start()

    time.sleep(5)
    stop.set()

    for t in threads:
        t.join(timeout=2)

    total_reader_ops = sum(st.ops for st in reader_stats)

    print(f"  Writer ops: {writer_stats.ops:,}")
    print(f"  Reader ops: {total_reader_ops:,}")
    print(f"  Errors: {len(errors)}")
    if errors:
        for e in errors[:5]:
            print(f"    {e}")

    return len(errors) == 0


def test_mask_table_ordering():
    """
    Test 6: Try to catch mask/table ordering issue.

    The concern: reader loads table, then mask. If resize happens between,
    reader has old table pointer but new mask value.

    In the patch:
    - Resize stores: table=NULL, then mask=new, then table=new
    - Reader loads: table, then mask, then checks table again

    The second table check should catch this, but let's stress it.
    """
    print("\n=== Test 6: Mask/table ordering stress ===")

    s = set()
    stop = threading.Event()
    barrier = threading.Barrier(9)  # 8 readers + 1 writer
    reader_stats = [Stats() for _ in range(8)]
    writer_stats = Stats()
    errors = []
    error_lock = threading.Lock()

    def reader(tid):
        stats = reader_stats[tid]
        barrier.wait()
        while not stop.is_set():
            try:
                # Very tight loop
                for i in range(1000):
                    _ = i % 50 in s
                    stats.ops += 1
            except Exception as e:
                with error_lock:
                    errors.append(f"Reader {tid}: {e}")
                stats.errors += 1

    def writer():
        barrier.wait()
        size = 10
        while not stop.is_set():
            try:
                # Alternate between small and large to force resize
                s.clear()
                s.update(range(size))
                writer_stats.ops += 2
                # Vary size to hit different resize points
                size = 10 if size > 50 else size + 10
            except Exception as e:
                with error_lock:
                    errors.append(f"Writer: {e}")

    threads = [threading.Thread(target=reader, args=(i,)) for i in range(8)]
    threads.append(threading.Thread(target=writer))

    for t in threads:
        t.start()

    time.sleep(5)
    stop.set()

    for t in threads:
        t.join(timeout=2)

    total_reader_ops = sum(st.ops for st in reader_stats)

    print(f"  Writer ops: {writer_stats.ops:,}")
    print(f"  Reader ops: {total_reader_ops:,}")
    print(f"  Errors: {len(errors)}")
    if errors:
        for e in errors[:5]:
            print(f"    {e}")

    return len(errors) == 0


def main():
    parser = argparse.ArgumentParser(description='Targeted resize race tests')
    parser.add_argument('--test', type=int, help='Run specific test (1-6)')
    args = parser.parse_args()

    tests = [
        ("Resize during contains (grow)", test_resize_during_contains),
        ("Clear during contains", test_clear_during_contains),
        ("Threshold crossing", test_threshold_crossing),
        ("Multiple writers clear", test_multiple_writers_clear),
        ("Rapid resize", test_rapid_resize),
        ("Mask/table ordering", test_mask_table_ordering),
    ]

    if args.test:
        if 1 <= args.test <= len(tests):
            name, func = tests[args.test - 1]
            result = func()
            print(f"\n{'PASSED' if result else 'FAILED'}: {name}")
            sys.exit(0 if result else 1)
        else:
            print(f"Invalid test number. Choose 1-{len(tests)}")
            sys.exit(1)

    # Run all tests
    results = []
    for name, func in tests:
        try:
            result = func()
            results.append((name, result))
        except Exception as e:
            print(f"  EXCEPTION: {e}")
            results.append((name, False))

    print("\n" + "=" * 50)
    print("SUMMARY")
    print("=" * 50)

    all_passed = True
    for name, passed in results:
        status = "PASSED" if passed else "FAILED"
        print(f"  {status}: {name}")
        if not passed:
            all_passed = False

    print("=" * 50)
    print(f"OVERALL: {'SUCCESS' if all_passed else 'FAILURE'}")
    sys.exit(0 if all_passed else 1)


if __name__ == '__main__':
    main()