[Perf] corehttp - updating async stream upload httpx test

sdk/core/azure-core/perf-tests.yml

@@ -14,16 +14,16 @@ Tests:
   Arguments:
   - --size 1024 --parallel 64 --duration 60 --policies all
   - --size 1024 --parallel 64 --duration 60 --policies all --use-entra-id
-  - --size 10240 --parallel 32 --duration 60
-  - --size 10240 --parallel 32 --duration 60 --transport requests
+  - --size 10240 --parallel 64 --duration 60
+  - --size 10240 --parallel 64 --duration 60 --transport requests
 
 - Test: download-binary
   Class: DownloadBinaryDataTest
   Arguments:
   - --size 1024 --parallel 64 --duration 60
   - --size 1024 --parallel 64 --duration 60 --transport requests
   - --size 1024 --parallel 64 --duration 60 --use-entra-id
-  - --size 10240 --parallel 32 --duration 60 --policies all
+  - --size 10240 --parallel 64 --duration 60 --policies all
 
 - Test: update-entity
   Class: UpdateEntityJSONTest

sdk/core/corehttp/perf-tests.yml

@@ -14,16 +14,16 @@ Tests:
   Arguments:
   - --size 1024 --parallel 64 --duration 60 --policies all
   - --size 1024 --parallel 64 --duration 60 --policies all --use-entra-id
-  - --size 10240 --parallel 32 --duration 60
-  - --size 10240 --parallel 32 --duration 60 --transport httpx
+  - --size 10240 --parallel 64 --duration 60
+  - --size 10240 --parallel 64 --duration 60 --transport httpx
 
 - Test: download-binary
   Class: DownloadBinaryDataTest
   Arguments:
   - --size 1024 --parallel 64 --duration 60
   - --size 1024 --parallel 64 --duration 60 --transport httpx
   - --size 1024 --parallel 64 --duration 60 --use-entra-id
-  - --size 10240 --parallel 32 --duration 60 --policies all
+  - --size 10240 --parallel 64 --duration 60 --policies all
 
 - Test: update-entity
   Class: UpdateEntityJSONTest

sdk/core/corehttp/tests/perf_tests/upload_binary.py

@@ -5,7 +5,7 @@
 
 from time import time
 from wsgiref.handlers import format_date_time
-from devtools_testutils.perfstress_tests import RandomStream, AsyncRandomStream
+from devtools_testutils.perfstress_tests import RandomStream, AsyncRandomStream, AsyncIteratorRandomStream
 
 from corehttp.rest import HttpRequest
 from corehttp.exceptions import (
@@ -29,7 +29,14 @@ def __init__(self, arguments):
         blob_name = "uploadtest"
         self.blob_endpoint = f"{self.account_endpoint}{self.container_name}/{blob_name}"
         self.upload_stream = RandomStream(self.args.size)
-        self.upload_stream_async = AsyncRandomStream(self.args.size)
+
+        # The AsyncIteratorRandomStream is used for upload stream scenario, since the
+        # async httpx transport requires the request body stream to be type AsyncIterator (i.e. have an __aiter__ and __anext__ method rather than __iter__).
+        # Specific check in httpx here: https://github.com/encode/httpx/blob/7df47ce4d93a06f2c3310cd692b4c2336d7663ba/httpx/_content.py#L116.
+        if self.args.transport == "httpx":
+            self.upload_stream_async = AsyncIteratorRandomStream(self.args.size)
+        else:
+            self.upload_stream_async = AsyncRandomStream(self.args.size)
 
     def run_sync(self):
         self.upload_stream.reset()

tools/azure-sdk-tools/devtools_testutils/perfstress_tests/__init__.py

@@ -9,7 +9,7 @@
 from ._perf_stress_runner import _PerfStressRunner
 from ._perf_stress_test import PerfStressTest
 from ._random_stream import RandomStream, WriteStream, get_random_bytes
-from ._async_random_stream import AsyncRandomStream
+from ._async_random_stream import AsyncRandomStream, AsyncIteratorRandomStream
 from ._batch_perf_test import BatchPerfTest
 from ._event_perf_test import EventPerfTest
 
@@ -19,6 +19,7 @@
     "EventPerfTest",
     "RandomStream",
     "WriteStream",
+    "AsyncIteratorRandomStream",
     "AsyncRandomStream",
     "get_random_bytes"
 ]

tools/azure-sdk-tools/devtools_testutils/perfstress_tests/_async_random_stream.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from typing import AsyncIterator
 from io import BytesIO
 
 from ._random_stream import get_random_bytes, _DEFAULT_LENGTH
@@ -58,3 +59,47 @@ def remaining(self):
 
     def close(self):
         self._closed = True
+
+
+class AsyncIteratorRandomStream(AsyncIterator[bytes]):
+    """
+        Async random stream of bytes for methods that accept AsyncIterator as input.
+    """
+    def __init__(self, length, initial_buffer_length=_DEFAULT_LENGTH):
+        super().__init__()
+        self._base_data = get_random_bytes(initial_buffer_length)
+        self._data_length = length
+        self._base_buffer_length = initial_buffer_length
+        self._position = 0
+        self._remaining = length
+
+    def __len__(self):
+        return self._remaining
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._remaining == 0:
+            raise StopAsyncIteration
+        return self.read()
+
+    def reset(self):
+        self._position = 0
+        self._remaining = self._data_length
+
+    def read(self, size=None):
+        if self._remaining == 0:
+            return b""
+
+        if size is None:
+            e = self._base_buffer_length
+        else:
+            e = size
+        e = min(e, self._remaining)
+        if e > self._base_buffer_length:
+            self._base_data = get_random_bytes(e)
+            self._base_buffer_length = e
+        self._remaining = self._remaining - e
+        self._position += e
+        return self._base_data[:e]