chore: optimization to reduce number of open TCP connections while running zonal system tests

cloudbuild/run_zonal_tests.sh

@@ -22,5 +22,6 @@ pip install -e .
 echo '--- Setting up environment variables on VM ---'
 export ZONAL_BUCKET=${_ZONAL_BUCKET}
 export RUN_ZONAL_SYSTEM_TESTS=True
-echo '--- Running Zonal tests on VM ---'
+CURRENT_ULIMIT=$(ulimit -n)
+echo '--- Running Zonal tests on VM with ulimit set to ---' $CURRENT_ULIMIT
 pytest -vv -s --log-format='%(asctime)s %(levelname)s %(message)s' --log-date-format='%H:%M:%S' tests/system/test_zonal.py

cloudbuild/zb-system-tests-cloudbuild.yaml

@@ -3,6 +3,7 @@ substitutions:
   _ZONE: "us-central1-a"
   _SHORT_BUILD_ID: ${BUILD_ID:0:8}
   _VM_NAME: "py-sdk-sys-test-${_SHORT_BUILD_ID}"
+  _ULIMIT: "10000" # 10k, for gRPC bidi streams
 
 
 
@@ -67,7 +68,7 @@ steps:
         # Execute the script on the VM via SSH.
         # Capture the exit code to ensure cleanup happens before the build fails.
         set +e
-        gcloud compute ssh ${_VM_NAME} --zone=${_ZONE} --internal-ip --ssh-key-file=/workspace/.ssh/google_compute_engine --command="COMMIT_SHA=${COMMIT_SHA} _ZONAL_BUCKET=${_ZONAL_BUCKET} bash run_zonal_tests.sh"
+        gcloud compute ssh ${_VM_NAME} --zone=${_ZONE} --internal-ip --ssh-key-file=/workspace/.ssh/google_compute_engine --command="ulimit -n {_ULIMIT}; COMMIT_SHA=${COMMIT_SHA} _ZONAL_BUCKET=${_ZONAL_BUCKET} bash run_zonal_tests.sh"
         EXIT_CODE=$?
         set -e
 

tests/system/test_zonal.py

@@ -1,5 +1,4 @@
 # py standard imports
-import asyncio
 import os
 import uuid
 from io import BytesIO
@@ -8,6 +7,7 @@
 import google_crc32c
 
 import pytest
+import gc
 
 # current library imports
 from google.cloud.storage._experimental.asyncio.async_grpc_client import AsyncGrpcClient
@@ -19,6 +19,7 @@
     AsyncMultiRangeDownloader,
 )
 
+
 pytestmark = pytest.mark.skipif(
     os.getenv("RUN_ZONAL_SYSTEM_TESTS") != "True",
     reason="Zonal system tests need to be explicitly enabled. This helps scheduling tests in Kokoro and Cloud Build.",
@@ -36,36 +37,6 @@ def _get_equal_dist(a: int, b: int) -> tuple[int, int]:
     return a + step, a + 2 * step
 
 
-async def write_one_appendable_object(
-    bucket_name: str,
-    object_name: str,
-    data: bytes,
-) -> None:
-    """Helper to write an appendable object."""
-    grpc_client = AsyncGrpcClient(attempt_direct_path=True).grpc_client
-    writer = AsyncAppendableObjectWriter(grpc_client, bucket_name, object_name)
-    await writer.open()
-    await writer.append(data)
-    await writer.close()
-
-
-@pytest.fixture(scope="function")
-def appendable_object(storage_client, blobs_to_delete):
-    """Fixture to create and cleanup an appendable object."""
-    object_name = f"appendable_obj_for_mrd-{str(uuid.uuid4())[:4]}"
-    asyncio.run(
-        write_one_appendable_object(
-            _ZONAL_BUCKET,
-            object_name,
-            _BYTES_TO_UPLOAD,
-        )
-    )
-    yield object_name
-
-    # Clean up; use json client (i.e. `storage_client` fixture) to delete.
-    blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
-
-
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "object_size",
@@ -114,6 +85,9 @@ async def test_basic_wrd(
 
     # Clean up; use json client (i.e. `storage_client` fixture) to delete.
     blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+    del writer
+    del mrd
+    gc.collect()
 
 
 @pytest.mark.asyncio
@@ -161,12 +135,20 @@ async def test_basic_wrd_in_slices(storage_client, blobs_to_delete, object_size)
 
     # Clean up; use json client (i.e. `storage_client` fixture) to delete.
     blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+    del writer
+    del mrd
+    gc.collect()
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "flush_interval",
-    [2 * 1024 * 1024, 4 * 1024 * 1024, 8 * 1024 * 1024, _DEFAULT_FLUSH_INTERVAL_BYTES],
+    [
+        2 * 1024 * 1024,
+        4 * 1024 * 1024,
+        8 * 1024 * 1024,
+        _DEFAULT_FLUSH_INTERVAL_BYTES,
+    ],
 )
 async def test_wrd_with_non_default_flush_interval(
     storage_client,
@@ -214,6 +196,9 @@ async def test_wrd_with_non_default_flush_interval(
 
     # Clean up; use json client (i.e. `storage_client` fixture) to delete.
     blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+    del writer
+    del mrd
+    gc.collect()
 
 
 @pytest.mark.asyncio
@@ -237,20 +222,28 @@ async def test_read_unfinalized_appendable_object(storage_client, blobs_to_delet
 
     # Clean up; use json client (i.e. `storage_client` fixture) to delete.
     blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+    del writer
+    del mrd
+    gc.collect()
 
 
 @pytest.mark.asyncio
-async def test_mrd_open_with_read_handle(appendable_object):
-    grpc_client = AsyncGrpcClient(attempt_direct_path=True).grpc_client
+async def test_mrd_open_with_read_handle():
+    grpc_client = AsyncGrpcClient().grpc_client
+    object_name = f"test_read_handl-{str(uuid.uuid4())[:4]}"
+    writer = AsyncAppendableObjectWriter(grpc_client, _ZONAL_BUCKET, object_name)
+    await writer.open()
+    await writer.append(_BYTES_TO_UPLOAD)
+    await writer.close()
 
-    mrd = AsyncMultiRangeDownloader(grpc_client, _ZONAL_BUCKET, appendable_object)
+    mrd = AsyncMultiRangeDownloader(grpc_client, _ZONAL_BUCKET, object_name)
     await mrd.open()
     read_handle = mrd.read_handle
     await mrd.close()
 
     # Open a new MRD using the `read_handle` obtained above
     new_mrd = AsyncMultiRangeDownloader(
-        grpc_client, _ZONAL_BUCKET, appendable_object, read_handle=read_handle
+        grpc_client, _ZONAL_BUCKET, object_name, read_handle=read_handle
     )
     await new_mrd.open()
     # persisted_size not set when opened with read_handle
@@ -259,3 +252,6 @@ async def test_mrd_open_with_read_handle(appendable_object):
     await new_mrd.download_ranges([(0, 0, buffer)])
     await new_mrd.close()
     assert buffer.getvalue() == _BYTES_TO_UPLOAD
+    del mrd
+    del new_mrd
+    gc.collect()