⚡️ Speed up method GrpcMockSysDB.reset_state by 15%

[codeflash-ai]

📄 15% (0.15x) speedup for GrpcMockSysDB.reset_state in chromadb/db/impl/grpc/server.py

⏱️ Runtime : 158 microseconds → 138 microseconds (best of 369 runs)

📝 Explanation and details

The optimization achieves a 14% speedup primarily through two key changes:

1. Pre-computed UUID constant: The original code called UUID(int=0) on every reset_state() call (taking 30% of execution time according to profiler). The optimized version uses a module-level _ZERO_UUID = UUID(int=0) constant, eliminating repeated UUID construction.

2. Dictionary initialization optimization: Instead of creating empty dictionaries and then populating them step-by-step with multiple dictionary lookups and assignments, the optimized version creates the nested dictionary structures in single statements:

• {tenant: {db: {}}} replaces 3 separate assignment operations
• {tenant: {db: _ZERO_UUID}} replaces 3 separate assignment operations

3. Local variable caching: Constants DEFAULT_TENANT and DEFAULT_DATABASE are cached in local variables to avoid repeated global lookups, though this provides minimal benefit.

The line profiler shows the original version spent significant time on the UUID construction (47,610ns) and multiple dictionary assignments. The optimized version reduces total execution time from 158μs to 138μs by eliminating the UUID construction overhead and reducing dictionary operation complexity.

These optimizations are most effective for test cases that frequently call reset_state(), as shown in the annotated tests where speedups range from 27% to over 100% depending on the scenario. The optimization maintains identical behavior while reducing computational overhead during state resets.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 39 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from uuid import UUID

# imports
import pytest
from chromadb.db.impl.grpc.server import GrpcMockSysDB

# Function to test: GrpcMockSysDB.reset_state
# We'll define minimal stubs for dependencies and the class itself for testing.

# Minimal stubs for dependencies
DEFAULT_TENANT = "default_tenant"
DEFAULT_DATABASE = "default_database"

class SystemStub:
    class SettingsStub:
        def require(self, key):
            if key == "chroma_server_grpc_port":
                return 12345
            raise KeyError(key)
    settings = SettingsStub()
from chromadb.db.impl.grpc.server import GrpcMockSysDB

# -------------------- UNIT TESTS --------------------

@pytest.fixture
def sysdb():
    # Fixture to create a fresh GrpcMockSysDB for each test
    return GrpcMockSysDB(SystemStub())

# ========== 1. BASIC TEST CASES ==========

def test_reset_state_clears_segments_and_collections(sysdb):
    """Basic: After reset_state, segments and collections are empty except for defaults."""
    # Setup: add some dummy data
    sysdb._segments['seg1'] = 'segment1'
    sysdb._tenants_to_databases_to_collections['foo'] = {'bar': {'baz': 'collection'}}
    sysdb._tenants_to_database_to_id['foo'] = {'bar': UUID(int=1)}
    # Call reset
    sysdb.reset_state() # 4.72μs -> 2.96μs (59.3% faster)

def test_reset_state_idempotent(sysdb):
    """Basic: Calling reset_state multiple times yields the same state."""
    sysdb._segments['seg'] = 'segment'
    sysdb.reset_state() # 5.27μs -> 2.62μs (101% faster)
    # Save state after first reset
    state1 = (
        dict(sysdb._segments),
        dict(sysdb._tenants_to_databases_to_collections),
        dict(sysdb._tenants_to_database_to_id),
    )
    sysdb.reset_state() # 2.50μs -> 1.26μs (98.5% faster)
    # Save state after second reset
    state2 = (
        dict(sysdb._segments),
        dict(sysdb._tenants_to_databases_to_collections),
        dict(sysdb._tenants_to_database_to_id),
    )

def test_reset_state_preserves_server_port(sysdb):
    """Basic: reset_state does not affect _server_port."""
    old_port = sysdb._server_port
    sysdb.reset_state() # 5.09μs -> 2.27μs (125% faster)

# ========== 2. EDGE TEST CASES ==========

def test_reset_state_with_no_prior_data(sysdb):
    """Edge: reset_state works if called on a fresh object (no prior data)."""
    sysdb.reset_state() # 4.95μs -> 2.33μs (113% faster)

def test_reset_state_with_non_string_keys(sysdb):
    """Edge: Handles non-string keys in dicts gracefully (should remove them)."""
    sysdb._segments[42] = 'numeric key'
    sysdb._tenants_to_databases_to_collections[None] = {'x': {'y': 'z'}}
    sysdb._tenants_to_database_to_id[3.14] = {'a': UUID(int=5)}
    sysdb.reset_state() # 3.53μs -> 2.23μs (58.3% faster)

def test_reset_state_overwrites_existing_defaults(sysdb):
    """Edge: If defaults exist with different data, they are reset."""
    sysdb._tenants_to_databases_to_collections[DEFAULT_TENANT] = {DEFAULT_DATABASE: {'foo': 'bar'}}
    sysdb._tenants_to_database_to_id[DEFAULT_TENANT] = {DEFAULT_DATABASE: UUID(int=999)}
    sysdb.reset_state() # 3.52μs -> 2.35μs (50.0% faster)

def test_reset_state_with_large_unexpected_data(sysdb):
    """Edge: Handles large, unexpected data in the dicts."""
    # Add large number of irrelevant tenants/databases/segments
    for i in range(100):
        sysdb._segments[f'seg{i}'] = f'segment{i}'
        sysdb._tenants_to_databases_to_collections[f'tenant{i}'] = {f'db{i}': {f'col{i}': f'collection{i}'}}
        sysdb._tenants_to_database_to_id[f'tenant{i}'] = {f'db{i}': UUID(int=i)}
    sysdb.reset_state() # 3.74μs -> 2.55μs (46.5% faster)

# ========== 3. LARGE SCALE TEST CASES ==========

def test_reset_state_large_scale(sysdb):
    """Large Scale: Handles 1000 tenants/databases/collections/segments."""
    # Fill with lots of data
    for i in range(1000):
        sysdb._segments[f'seg{i}'] = f'segment{i}'
        t = f'tenant{i}'
        d = f'db{i}'
        c = f'col{i}'
        sysdb._tenants_to_databases_to_collections.setdefault(t, {})[d] = {c: f'collection{i}'}
        sysdb._tenants_to_database_to_id.setdefault(t, {})[d] = UUID(int=i)
    sysdb.reset_state() # 4.96μs -> 3.40μs (46.2% faster)

def test_reset_state_performance_large_scale(sysdb):
    """Large Scale: reset_state runs efficiently on large data (not a timing test, but checks correctness)."""
    # Add 999 segments and tenants (under the 1000 limit)
    for i in range(999):
        sysdb._segments[f'seg{i}'] = f'segment{i}'
        sysdb._tenants_to_databases_to_collections[f'tenant{i}'] = {f'db{i}': {f'col{i}': f'collection{i}'}}
        sysdb._tenants_to_database_to_id[f'tenant{i}'] = {f'db{i}': UUID(int=i)}
    sysdb.reset_state() # 4.71μs -> 3.71μs (27.1% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from uuid import UUID

# imports
import pytest
from chromadb.db.impl.grpc.server import GrpcMockSysDB

# Function to test: a minimal, self-contained reset_state for a mock component.
# This version matches the behavior described in chromadb/db/impl/grpc/server.py above.
DEFAULT_TENANT = "default"
DEFAULT_DATABASE = "default"

class MockComponent:
    def __init__(self):
        self._segments = {"seg1": "val1", "seg2": "val2"}
        self._collection_to_segments = {"coll1": ["seg1"], "coll2": ["seg2"]}
        self._tenants_to_databases_to_collections = {
            "tenant1": {"db1": {"collA": "A"}, "db2": {"collB": "B"}},
            "tenant2": {"db3": {"collC": "C"}},
        }
        self._tenants_to_database_to_id = {
            "tenant1": {"db1": UUID(int=1), "db2": UUID(int=2)},
            "tenant2": {"db3": UUID(int=3)},
        }

    def reset_state(self):
        # Reset all state to initial blank/default state
        self._segments = {}
        self._tenants_to_databases_to_collections = {}
        # Create defaults
        self._tenants_to_databases_to_collections[DEFAULT_TENANT] = {}
        self._tenants_to_databases_to_collections[DEFAULT_TENANT][DEFAULT_DATABASE] = {}
        self._tenants_to_database_to_id = {}
        self._tenants_to_database_to_id[DEFAULT_TENANT] = {}
        self._tenants_to_database_to_id[DEFAULT_TENANT][DEFAULT_DATABASE] = UUID(int=0)

# =========================
# Basic Test Cases
# =========================

def test_reset_state_basic_clears_segments():
    """Basic: After reset_state, _segments should be empty dict."""
    comp = MockComponent()
    comp._segments["extra"] = "value"
    comp.reset_state() # 2.48μs -> 2.40μs (3.37% faster)

def test_reset_state_basic_clears_collections():
    """Basic: After reset_state, _tenants_to_databases_to_collections should have only default tenant/database."""
    comp = MockComponent()
    comp._tenants_to_databases_to_collections["other"] = {"db": {"coll": "val"}}
    comp.reset_state() # 2.20μs -> 2.17μs (1.57% faster)

def test_reset_state_basic_sets_default_database_id():
    """Basic: After reset_state, _tenants_to_database_to_id should have default tenant/database with UUID(int=0)."""
    comp = MockComponent()
    comp._tenants_to_database_to_id["other"] = {"db": UUID(int=999)}
    comp.reset_state() # 2.16μs -> 2.08μs (3.70% faster)

def test_reset_state_basic_clears_collection_to_segments():
    """Basic: _collection_to_segments should not be reset by reset_state (not in function)."""
    comp = MockComponent()
    original = comp._collection_to_segments.copy()
    comp.reset_state() # 2.06μs -> 2.05μs (0.733% faster)

# =========================
# Edge Test Cases
# =========================

def test_reset_state_edge_empty_state():
    """Edge: Resetting already-empty state should still create defaults."""
    comp = MockComponent()
    comp._segments = {}
    comp._tenants_to_databases_to_collections = {}
    comp._tenants_to_database_to_id = {}
    comp.reset_state() # 1.73μs -> 1.69μs (2.42% faster)

def test_reset_state_edge_mutation_protection():
    """Edge: Changing defaults after reset_state does not affect future resets."""
    comp = MockComponent()
    comp.reset_state() # 2.10μs -> 2.11μs (0.521% slower)
    # Mutate the defaults
    comp._tenants_to_databases_to_collections[DEFAULT_TENANT][DEFAULT_DATABASE]["collX"] = "X"
    comp._tenants_to_database_to_id[DEFAULT_TENANT][DEFAULT_DATABASE] = UUID(int=123)
    # Reset again and check that mutation is wiped
    comp.reset_state() # 1.27μs -> 1.34μs (5.45% slower)

def test_reset_state_edge_non_string_keys():
    """Edge: State with non-string keys is wiped by reset_state."""
    comp = MockComponent()
    comp._tenants_to_databases_to_collections[42] = {"db": {"coll": "val"}}
    comp._tenants_to_database_to_id[None] = {"db": UUID(int=5)}
    comp.reset_state() # 2.12μs -> 2.10μs (1.38% faster)


def test_reset_state_large_scale_many_tenants_and_dbs():
    """Large Scale: Reset after many tenants/databases/collections."""
    comp = MockComponent()
    # Fill with many tenants, dbs, collections
    for i in range(1000):
        tenant = f"tenant{i}"
        db = f"db{i}"
        coll = f"coll{i}"
        comp._tenants_to_databases_to_collections.setdefault(tenant, {})[db] = {coll: f"value{i}"}
        comp._tenants_to_database_to_id.setdefault(tenant, {})[db] = UUID(int=i)
    comp.reset_state() # 86.4μs -> 83.9μs (3.04% faster)

def test_reset_state_large_scale_segments():
    """Large Scale: Reset after many segments."""
    comp = MockComponent()
    # Fill with many segments
    for i in range(1000):
        comp._segments[f"seg{i}"] = f"value{i}"
    comp.reset_state() # 9.83μs -> 9.94μs (1.02% slower)

def test_reset_state_large_scale_collection_to_segments():
    """Large Scale: _collection_to_segments is not affected by reset_state."""
    comp = MockComponent()
    comp._collection_to_segments = {f"coll{i}": [f"seg{i}"] for i in range(1000)}
    original = comp._collection_to_segments.copy()
    comp.reset_state() # 2.40μs -> 2.34μs (2.57% faster)

# =========================
# Mutation Testing: Negative Test
# =========================

To edit these changes git checkout codeflash/optimize-GrpcMockSysDB.reset_state-mh1lhw8e and push.