Change Workloads For Circuit Breaker (#40843)

* change workloads based on feedback

* Update sdk/cosmos/azure-cosmos/tests/workloads/dev.md

Co-authored-by: Copilot <[email protected]>

* Update sdk/cosmos/azure-cosmos/tests/workloads/workload_utils.py

Co-authored-by: Copilot <[email protected]>

* react to comments

---------

Co-authored-by: Copilot <[email protected]>

Author: tvaron3

sdk/cosmos/azure-cosmos/tests/workloads/dev.md

@@ -3,12 +3,13 @@ This directory contains the scale testing workloads for the SDK. The workloads a
 and scalability of the SDK under various conditions. There are different types of workloads and each will create a log 
 file when run. These logs are named in this format `<file name>-<process id>-<date time>.log`. 
 
-### Setup VM
+### Setup Scale Testing
 1. Create a VM in Azure with the following configuration:
    - 8 vCPUs
    - 32 GB RAM
    - Ubuntu
    - Accelerated networking
+1. Create an Azure App Insights Resource (Optional)
 1. Fork and clone this repository
 1. Go to azure cosmos folder
    - `cd azure-sdk-for-python/sdk/cosmos/azure-cosmos`
@@ -19,6 +20,7 @@ file when run. These logs are named in this format `<file name>-<process id>-<da
    - `python3 -m venv azure-cosmosdb-sdk-environment`
    - `source azure-cosmosdb-sdk-environment/bin/activate`
    - `pip install -r dev_requirements.txt`
+   - `pip install azure-monitor-opentelemetry`
 1. Checkout the branch with the changes to test. 
 1. Install azure-cosmos
    - `pip install .`
@@ -40,3 +42,9 @@ file when run. These logs are named in this format `<file name>-<process id>-<da
 ### Monitor Run
 - `ps -eaf | grep "python3"` to see the running processes
 - `tail -f <log_file>` to see the logs in real time 
+
+### Close Workloads
+- If you want to keep the logs and stop the scripts,  
+   `./shutdown_workloads.sh --no-remove-logs` 
+- If you want to remove the logs and stop the scripts,        
+   `./shutdown_workloads.sh`

sdk/cosmos/azure-cosmos/tests/workloads/r_w_q_workload.py

@@ -1,6 +1,5 @@
 # The MIT License (MIT)
 # Copyright (c) Microsoft Corporation. All rights reserved.
-import os
 import sys
 
 from azure.cosmos import documents

sdk/cosmos/azure-cosmos/tests/workloads/shutdown_workloads.sh

@@ -1,8 +1,28 @@
 #!/bin/bash
 # cspell:disable
+
+NO_REMOVE_LOGS=false
+
+while [[ "$#" -gt 0 ]]; do
+    case $1 in
+        --no-remove-logs)
+            NO_REMOVE_LOGS=true
+            ;;
+        *)
+            echo "Unknown parameter: $1"
+            exit 1
+            ;;
+    esac
+    shift
+done
+
 # Kill all python3 processes
 pkill -f "python3 ./r_"
 pkill -f "python3 ./w_"
 
-# Remove all files with .logs in their filename
-rm -f *.log*
+if [ "$NO_REMOVE_LOGS" = false ]; then
+    # Remove all files with .log in their filename
+    rm -f *.log*
+    # Remove envoy log files
+    rm envoy/logs/*
+fi

sdk/cosmos/azure-cosmos/tests/workloads/workload_configs.py

@@ -1,6 +1,8 @@
 # The MIT License (MIT)
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Replace with your Cosmos DB details
+import logging
+
 PREFERRED_LOCATIONS = []
 CLIENT_EXCLUDED_LOCATIONS = []
 REQUEST_EXCLUDED_LOCATIONS = []
@@ -9,6 +11,10 @@
 COSMOS_KEY = ""
 COSMOS_CONTAINER = "scale_cont"
 COSMOS_DATABASE = "scale_db"
+USER_AGENT_PREFIX = ""
+LOG_LEVEL = logging.DEBUG
+APP_INSIGHTS_CONNECTION_STRING = ""
+CIRCUIT_BREAKER_ENABLED = False
 USE_MULTIPLE_WRITABLE_LOCATIONS = False
 CONCURRENT_REQUESTS = 100
 CONCURRENT_QUERIES = 2

sdk/cosmos/azure-cosmos/tests/workloads/workload_utils.py

@@ -1,19 +1,21 @@
 # The MIT License (MIT)
 # Copyright (c) Microsoft Corporation. All rights reserved.
 import asyncio
-import logging
 import os
 import random
 from datetime import datetime
 from logging.handlers import RotatingFileHandler
-from workload_configs import NUMBER_OF_LOGICAL_PARTITIONS, PARTITION_KEY
+
+from azure.monitor.opentelemetry import configure_azure_monitor
+from workload_configs import *
 
 _NOISY_ERRORS = set([404, 409, 412])
 _NOISY_SUB_STATUS_CODES = set([0, None])
 _REQUIRED_ATTRIBUTES = ["resource_type", "verb", "operation_type", "status_code", "sub_status_code", "duration"]
 
 def get_user_agent(client_id):
-    return str(client_id) + "-" + datetime.now().strftime("%Y%m%d-%H%M%S")
+    prefix = USER_AGENT_PREFIX + "-" if USER_AGENT_PREFIX else ""
+    return prefix + str(client_id) + "-" + datetime.now().strftime("%Y%m%d-%H%M%S")
 
 def get_random_item():
     random_int = random.randint(0, NUMBER_OF_LOGICAL_PARTITIONS)
@@ -103,15 +105,21 @@ async def perform_query_concurrently(container, excluded_locations):
     items = [item async for item in results]
 
 def create_logger(file_name):
-    logger = logging.getLogger('azure.cosmos')
+    os.environ["AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"] = str(CIRCUIT_BREAKER_ENABLED)
+    logger = logging.getLogger()
+    if APP_INSIGHTS_CONNECTION_STRING:
+        configure_azure_monitor(
+            logger_name="azure.cosmos",
+            connection_string=APP_INSIGHTS_CONNECTION_STRING,
+        )
     prefix = os.path.splitext(file_name)[0] + "-" + str(os.getpid())
     # Create a rotating file handler
     handler = RotatingFileHandler(
-        "log-" + prefix + "-" + datetime.now().strftime("%Y%m%d-%H%M%S") + '.log',
+        "log-" + get_user_agent(prefix) + '.log',
         maxBytes=1024 * 1024 * 10,  # 10 mb
-        backupCount=3
+        backupCount=2
     )
-    logger.setLevel(logging.DEBUG)
+    logger.setLevel(LOG_LEVEL)
     # create filters for the logger handler to reduce the noise
     workload_logger_filter = WorkloadLoggerFilter()
     handler.addFilter(workload_logger_filter)