Add support for gRPC transport

osbenchmark/benchmark.py

@@ -587,6 +587,11 @@ def add_workload_source(subparser):
         "--worker-ips",
         help="Define a comma-separated list of hosts which should generate load (default: localhost).",
         default="localhost")
+    test_run_parser.add_argument(
+        "--grpc-target-hosts",
+        help="Define a comma-separated list of host:port pairs for gRPC endpoints "
+             "(default: localhost:9400).",
+        default="")
     test_run_parser.add_argument(
         "--client-options",
         "-c",
@@ -1070,6 +1075,10 @@ def configure_connection_params(arg_parser, args, cfg):
     cfg.add(config.Scope.applicationOverride, "client", "hosts", target_hosts)
     client_options = opts.ClientOptions(args.client_options, target_hosts=target_hosts)
     cfg.add(config.Scope.applicationOverride, "client", "options", client_options)
+
+    # Configure gRPC target hosts
+    grpc_target_hosts = opts.TargetHosts(args.grpc_target_hosts) if args.grpc_target_hosts else None
+    cfg.add(config.Scope.applicationOverride, "client", "grpc_hosts", grpc_target_hosts)
     if "timeout" not in client_options.default:
         console.info("You did not provide an explicit timeout in the client options. Assuming default of 10 seconds.")
     if list(target_hosts.all_hosts) != list(client_options.all_client_options):

osbenchmark/client.py

@@ -30,13 +30,16 @@
 import urllib3
 from urllib3.util.ssl_ import is_ipaddress
 
+import grpc
+from opensearch.protobufs.services.document_service_pb2_grpc import DocumentServiceStub
+from opensearch.protobufs.services.search_service_pb2_grpc import SearchServiceStub
+
 from osbenchmark.kafka_client import KafkaMessageProducer
 from osbenchmark import exceptions, doc_link, async_connection
 from osbenchmark.context import RequestContextHolder
 from osbenchmark.utils import console, convert
 from osbenchmark.cloud_provider import CloudProviderFactory
 
-
 class OsClientFactory:
     """
     Abstracts how the OpenSearch client is created. Intended for testing.
@@ -175,9 +178,11 @@ def create_async(self):
         # pylint: disable=import-outside-toplevel
         import io
         import aiohttp
-
         from opensearchpy.serializer import JSONSerializer
 
+        class BenchmarkAsyncOpenSearch(opensearchpy.AsyncOpenSearch, RequestContextHolder):
+            pass
+
         class LazyJSONSerializer(JSONSerializer):
             def loads(self, s):
                 meta = BenchmarkAsyncOpenSearch.request_context.get()
@@ -202,9 +207,6 @@ async def on_request_end(session, trace_config_ctx, params):
         self.client_options["serializer"] = LazyJSONSerializer()
         self.client_options["trace_config"] = trace_config
 
-        class BenchmarkAsyncOpenSearch(opensearchpy.AsyncOpenSearch, RequestContextHolder):
-            pass
-
         if self.provider:
             self.logger.info("Creating OpenSearch Async Client with provider %s", self.provider)
             return self.provider.create_client(self.hosts, self.client_options,
@@ -278,3 +280,124 @@ async def create(params):
             return await KafkaMessageProducer.create(params)
         else:
             raise ValueError(f"Unsupported ingestion source type: {producer_type}")
+
+
+class GrpcClientFactory:
+    """
+    Factory for creating gRPC client stubs.
+    Note gRPC channels must default `use_local_subchannel_pool` to true.
+    Sub channels manage the underlying connection with the server. When the global sub channel pool is used gRPC will
+    re-use sub channels and their underlying connections which does not appropriately reflect a multi client scenario.
+    """
+    def __init__(self, grpc_hosts):
+        self.grpc_hosts = grpc_hosts
+        self.logger = logging.getLogger(__name__)
+        self.grpc_channel_options = [
+            ('grpc.use_local_subchannel_pool', 1),
+            ('grpc.max_send_message_length', 10 * 1024 * 1024),  # 10 MB
+            ('grpc.max_receive_message_length', 10 * 1024 * 1024)  # 10 MB
+        ]
+
+    def create_grpc_stubs(self):
+        """
+        Create gRPC service stubs.
+        Returns a dict of {cluster_name: {service_name: stub}} structure.
+        """
+        stubs = {}
+
+        if len(self.grpc_hosts.all_hosts.items()) > 1:
+            raise NotImplementedError("Only one gRPC cluster is supported.")
+
+        if len(self.grpc_hosts.all_hosts["default"]) > 1:
+            raise NotImplementedError("Only one gRPC host is supported.")
+
+        host = self.grpc_hosts.all_hosts["default"][0]
+        grpc_addr = f"{host['host']}:{host['port']}"
+
+        self.logger.info("Creating gRPC channel for cluster default cluster at %s", grpc_addr)
+        channel = grpc.aio.insecure_channel(
+            target=grpc_addr,
+            options=self.grpc_channel_options,
+            compression=None
+        )
+
+        # Retain a reference to underlying channel in our stubs dictionary for graceful shutdown.
+        stubs["default"] = {
+            'document_service': DocumentServiceStub(channel),
+            'search_service': SearchServiceStub(channel),
+            '_channel': channel
+        }
+
+        return stubs
+
+
+class UnifiedClient:
+    """
+    Unified client that wraps both OpenSearch REST client and gRPC stubs.
+    This provides a single interface for runners to access both protocols.
+    Acts as a transparent proxy to the OpenSearch client while adding gRPC capabilities.
+    """
+    def __init__(self, opensearch_client, grpc_stubs=None):
+        self._opensearch = opensearch_client
+        self._grpc_stubs = grpc_stubs
+        self._logger = logging.getLogger(__name__)
+
+    def __getattr__(self, name):
+        """Delegate all unknown attributes to the underlying OpenSearch client."""
+        return getattr(self._opensearch, name)
+
+    def document_service(self, cluster_name="default"):
+        """Get the gRPC DocumentService stub for the specified cluster."""
+        if cluster_name in self._grpc_stubs:
+            return self._grpc_stubs[cluster_name].get('document_service')
+        else:
+            raise exceptions.SystemSetupError(
+                "gRPC DocumentService not available. Please configure --grpc-target-hosts.")
+
+    def search_service(self, cluster_name="default"):
+        """Get the gRPC SearchService stub for the specified cluster."""
+        if cluster_name in self._grpc_stubs:
+            return self._grpc_stubs[cluster_name].get('search_service')
+        else:
+            raise exceptions.SystemSetupError(
+                "gRPC SearchService not available. Please configure --grpc-target-hosts.")
+
+    def __del__(self):
+        """Close all gRPC channels."""
+        for cluster_stubs in self._grpc_stubs.values():
+            if '_channel' in cluster_stubs:
+                try:
+                    cluster_stubs['_channel'].close()
+                except Exception as e:
+                    self._logger.warning("Error closing gRPC channel: %s", e)
+        self._opensearch.close()
+
+    @property
+    def opensearch(self):
+        """Provide access to the underlying OpenSearch client for explicit access."""
+        return self._opensearch
+
+
+class UnifiedClientFactory:
+    """
+    Factory that creates UnifiedClient instances with both REST and gRPC support.
+    """
+    def __init__(self, rest_client_factory, grpc_hosts=None):
+        self.rest_client_factory = rest_client_factory
+        self.grpc_hosts = grpc_hosts
+        self.logger = logging.getLogger(__name__)
+
+    def create(self):
+        """Non async client is deprecated."""
+        raise NotImplementedError()
+
+    def create_async(self):
+        """Create a UnifiedClient with async REST client."""
+        opensearch_client = self.rest_client_factory.create_async()
+        grpc_stubs = None
+
+        if self.grpc_hosts:
+            grpc_factory = GrpcClientFactory(self.grpc_hosts)
+            grpc_stubs = grpc_factory.create_grpc_stubs()
+
+        return UnifiedClient(opensearch_client, grpc_stubs)

osbenchmark/worker_coordinator/proto_helpers/ProtoBulkHelper.py

@@ -0,0 +1,71 @@
+from opensearch.protobufs.schemas import document_pb2
+
+def _parse_docs_from_body(body):
+    index_op_lines = body.decode('utf-8').split('\n')
+    doc_list = []
+    for doc in index_op_lines[1::2]:
+        doc_list.append(doc)
+    return doc_list
+
+class ProtoBulkHelper:
+    # Build protobuf SearchRequest.
+    # Consumed from params dictionary:
+    # * ``body``: JSON body of bulk ingest request
+    # * ``index``: index name
+    @staticmethod
+    def build_proto_request(params):
+        index = params.get("index")
+        body = params.get("body")
+        doc_list = _parse_docs_from_body(body)
+        request = document_pb2.BulkRequest()
+        request.index = index
+        # All bulk requests here are index ops
+        op_container = document_pb2.OperationContainer()
+        op_container.index.CopyFrom(document_pb2.IndexOperation())
+        for doc in doc_list:
+            request_body = document_pb2.BulkRequestBody()
+            request_body.object = doc.encode('utf-8')
+            request_body.operation_container.CopyFrom(op_container)
+            request.request_body.append(request_body)
+        return request
+
+    # Parse stats from protobuf response.
+    # Consumed from params dictionary:
+    # ``index``: index name
+    # ``bulk-size``: documents per bulk request
+    # ``unit``: in the case of bulk always 'ops'
+    # ``detailed-results``: gRPC/Protobuf does not support detailed results at this time.
+    @staticmethod
+    def build_stats(response : document_pb2.BulkResponse, params):
+        if params.get("detailed-results"):
+            raise Exception("Detailed results not supported for gRPC bulk requests")
+
+        took = None
+        error_count = 0
+        success_count = 0
+        if response.errors:
+            error_count = params.get("bulk-size")
+        else:
+            took = response.took
+            for item in response.items:
+                # status field mirrors http code conventions
+                # https://github.com/opensearch-project/opensearch-protobufs/blob/b6f889416da83b7dc4a0408347965e7820bd61d0/protos/schemas/document.proto#L217-L219
+                if item.index.status > 299:
+                    error_count += 1
+                else:
+                    success_count += 1
+
+        meta_data = {
+            "index": params.get("index"),
+            "weight": params.get("bulk-size"),
+            "unit": params.get("unit"),
+            "took": took,
+            "success": error_count == 0,
+            "success-count": success_count,
+            "error-count": error_count,
+        }
+
+        if error_count > 0:
+            meta_data["error-type"] = "bulk"
+
+        return meta_data