Add mongo ycsb driver support

Change-Id: I9df090e623d386b72363d46ae422204e99d930b8
Reviewed-on: https://review.couchbase.org/c/perfrunner/+/245470
Reviewed-by: Salim Salim <salim.salim@couchbase.com>
Tested-by: Build Bot <build@couchbase.com>

Author: DragosPetruTaraban

perfrunner/helpers/local.py

@@ -930,6 +930,146 @@ def run_ycsb(
         run_local_shell_command(cmd_args, quiet=True, cwd=cwd, stderr=stderr_file)
 
 
+def run_mongo_ycsb(
+    hosts: List[str],
+    database: str,
+    action: str,
+    workload: str,
+    items: int,
+    workers: int,
+    target: int,
+    port: int = 27017,
+    instance: int = 0,
+    username: Optional[str] = None,
+    password: Optional[str] = None,
+    replica_set_name: Optional[str] = None,
+    ycsb_client: str = "mongodb",
+    write_concern: str = "majority",
+    batch_size: int = 1,
+    requestdistribution: str = "zipfian",
+    fieldlength: int = 1024,
+    fieldcount: int = 10,
+    upsert: bool = False,
+    write_all_fields: bool = True,
+    journal_ack: bool = True,
+    read_preference: Optional[str] = None,
+    read_concern: Optional[str] = None,
+    max_pool_size: Optional[int] = None,
+    ops: Optional[int] = None,
+    execution_time: Optional[int] = None,
+    timeseries: int = 0,
+    phase_params: Optional[dict] = None,
+    insert_test_params: Optional[dict] = None,
+    soe_params: Optional[dict] = None,
+    ycsb_jvm_args: Optional[str] = None,
+    measurement_type: Optional[str] = None,
+    histogram_buckets: Optional[int] = None,
+    histogram_bucket_size: Optional[int] = None,
+    verbose_histogram: bool = False,
+    ycsb_status_output: bool = True,
+):
+    seeds = ",".join(f"{h}:{port}" for h in hosts)
+
+    cred = ""
+    if username and password:
+        cred = (
+            f"{urllib.parse.quote(username, safe='')}:"
+            f"{urllib.parse.quote(password, safe='')}@"
+        )
+
+    uri_params = [f"w={write_concern}"]
+    if replica_set_name:
+        uri_params.append(f"replicaSet={replica_set_name}")
+    if not journal_ack:
+        uri_params.append("journal=false")
+    if read_preference:
+        uri_params.append(f"readPreference={read_preference}")
+    if read_concern:
+        uri_params.append(f"readConcernLevel={read_concern}")
+    if max_pool_size:
+        uri_params.append(f"maxPoolSize={max_pool_size}")
+    mongo_url = f"mongodb://{cred}{seeds}/{database}?" + "&".join(uri_params)
+
+    parameters = [
+        f"target={target}",
+        f"fieldlength={fieldlength}",
+        f"fieldcount={fieldcount}",
+        f"requestdistribution={requestdistribution}",
+        f"writeallfields={str(write_all_fields).lower()}",
+        f"mongodb.url={mongo_url}",
+        f"mongodb.database={database}",
+        f"mongodb.batchsize={batch_size}",
+        f"mongodb.writeConcern={write_concern}",
+        f"mongodb.upsert={str(upsert).lower()}",
+        f"exportfile=ycsb_{action}_{instance}_{database}.log",
+    ]
+
+    if ops is not None:
+        parameters += [f"operationcount={ops}"]
+    if execution_time is not None:
+        parameters += [f"maxexecutiontime={execution_time}"]
+    if ycsb_jvm_args is not None:
+        parameters += [f"jvm-args='{ycsb_jvm_args}'"]
+
+    if timeseries:
+        parameters += ["measurementtype=timeseries"]
+    elif measurement_type:
+        parameters += [f"measurementtype={measurement_type}"]
+    if histogram_buckets:
+        parameters += [f"histogram.buckets={histogram_buckets}"]
+    if histogram_bucket_size:
+        parameters += [f"histogram.bucket.size={histogram_bucket_size}"]
+    if verbose_histogram:
+        parameters += ["measurement.histogram.verbose=true"]
+
+    if soe_params:
+        parameters += [
+            f"totalrecordcount={items}",
+            f"recordcount={soe_params['recorded_load_cache_size']}",
+            f"insertstart={soe_params['insertstart']}",
+        ]
+    elif phase_params:
+        parameters += [
+            f"totalrecordcount={items}",
+            f"recordcount={phase_params['inserts_per_workerinstance']}",
+            f"insertstart={phase_params['insertstart']}",
+        ]
+    elif insert_test_params:
+        parameters += [
+            f"recordcount={insert_test_params['recordcount']}",
+            f"insertstart={insert_test_params['insertstart']}",
+        ]
+    else:
+        parameters += [f"recordcount={items}"]
+
+    cwd = "YCSB"
+    max_arg_strlen = get_max_arg_strlen()
+    if too_long_params := [arg for arg in parameters if len(arg) > max_arg_strlen]:
+        parameters = [arg for arg in parameters if arg not in too_long_params]
+        logger.info(f"Some parameters are too long to provide via CLI: {too_long_params}")
+        custom_workload = f"{workload}_custom_{uuid4().hex[:6]}"
+        logger.info(
+            f"Copying {workload} file to {custom_workload} and inserting long parameters there."
+        )
+        inject_params_into_ycsb_workload_file(
+            f"{cwd}/{workload}", f"{cwd}/{custom_workload}", too_long_params
+        )
+        workload = custom_workload
+
+    cmd_args = ["bin/ycsb", action, ycsb_client]
+    if ycsb_status_output:
+        cmd_args.append("-s")
+    cmd_args += ["-threads", workers, "-P", workload] + [
+        arg for param in parameters for arg in ["-p", param]
+    ]
+    cmd_args = list(map(str, cmd_args))
+
+    logger.info(f"Running: {shlex.join(cmd_args)}")
+    run_local_shell_command("pyenv local 2.7.18", quiet=True, cwd=cwd)
+    with open(f"{cwd}/ycsb_{action}_{instance}_{database}_stderr.log", "w") as stderr_file:
+        run_local_shell_command(cmd_args, quiet=True, cwd=cwd, stderr=stderr_file)
+
+
 def inject_params_into_ycsb_workload_file(
     source_filename: str, target_filename: str, params: list[str]
 ):

perfrunner/helpers/worker.py

@@ -60,7 +60,12 @@
 )
 from perfrunner.workloads.vectordb_bench import run_vectordb_bench_case
 from perfrunner.workloads.xdcr_conflict_sim import run_conflictsim
-from perfrunner.workloads.ycsb import ycsb_data_load, ycsb_workload
+from perfrunner.workloads.ycsb import (
+    ycsb_data_load,
+    ycsb_mongo_data_load,
+    ycsb_mongo_workload,
+    ycsb_workload,
+)
 
 try:
     set_start_method("fork")
@@ -144,6 +149,16 @@ def ycsb_data_load_task(*args):
     ycsb_data_load(*args)
 
 
+@celery.task
+def ycsb_mongo_data_load_task(*args):
+    ycsb_mongo_data_load(*args)
+
+
+@celery.task
+def ycsb_mongo_workload_task(*args):
+    ycsb_mongo_workload(*args)
+
+
 @celery.task
 def ycsb_task(*args):
     ycsb_workload(*args)

perfrunner/settings.py

@@ -1808,6 +1808,9 @@ def __init__(self, options: dict):
         self.ycsb_client = options.get('ycsb_client', self.YCSB_CLIENT)
         self.ycsb_out_of_order = int(options.get('out_of_order', self.YCSB_OUT_OF_ORDER))
         self.insertstart = int(options.get('insertstart', self.YCSB_INSERTSTART))
+        self.upsert = maybe_atoi(options.get("upsert", "false"))
+        self.mongo_batch_size = options.get("mongo_batch_size")
+        self.mongo_durability_level = options.get("mongo_durability_level")
         self.ycsb_split_workload = int(options.get('ycsb_split_workload', self.YCSB_SPLIT_WORKLOAD))
 
         self.range_scan_sampling = options.get('range_scan_sampling', self.RANGE_SCAN_SAMPLING)

perfrunner/workloads/ycsb.py

@@ -1,4 +1,4 @@
-from perfrunner.helpers.local import restart_memcached, run_ycsb
+from perfrunner.helpers.local import restart_memcached, run_mongo_ycsb, run_ycsb
 from perfrunner.settings import PhaseSettings, TargetSettings
 
 
@@ -148,3 +148,98 @@ def ycsb_workload(workload_settings: PhaseSettings,
         histogram_bucket_size=workload_settings.histogram_bucket_size,
         verbose_histogram=workload_settings.verbose_histogram,
     )
+
+def ycsb_mongo_data_load(workload_settings: PhaseSettings,
+                         target: TargetSettings,
+                         timer: int,
+                         instance: int):
+    phase_params = None
+    if workload_settings.phase:
+        phase_params = {
+            'insertstart': instance * workload_settings.inserts_per_workerinstance +
+            workload_settings.insertstart,
+            'inserts_per_workerinstance': workload_settings.inserts_per_workerinstance,
+        }
+
+    host = target.node
+    if target.cloud:
+        host = target.cloud.get("cluster_svc", host)
+
+    run_mongo_ycsb(
+        hosts=[host],
+        database=target.bucket,
+        username=target.username,
+        password=target.password,
+        action="load",
+        ycsb_client=workload_settings.ycsb_client,
+        workload=workload_settings.workload_path,
+        items=int(workload_settings.items),
+        workers=workload_settings.workers,
+        target=int(workload_settings.target),
+        instance=instance,
+        fieldlength=workload_settings.field_length,
+        fieldcount=workload_settings.field_count,
+        upsert=bool(workload_settings.upsert),
+        batch_size=workload_settings.mongo_batch_size or 1,
+        write_concern=workload_settings.mongo_durability_level or "majority",
+        requestdistribution=workload_settings.requestdistribution,
+        ycsb_jvm_args=workload_settings.ycsb_jvm_args,
+        timeseries=workload_settings.timeseries,
+        phase_params=phase_params,
+    )
+
+
+def ycsb_mongo_workload(workload_settings: PhaseSettings,
+                        target: TargetSettings,
+                        timer: int,
+                        instance: int):
+
+    if workload_settings.ycsb_split_workload:
+        split_instance = workload_settings.workload_instances // 2
+
+        if instance < split_instance:
+            workload_settings.workload_path = workload_settings.workload_path.split(",")[0]
+        elif instance >= split_instance:
+            workload_settings.workload_path = workload_settings.workload_path.split(",")[1]
+
+    insert_test_params = None
+    if workload_settings.insert_test_flag:
+        insert_test_params = {
+            'insertstart': int(instance * workload_settings.inserts_per_workerinstance +
+                               workload_settings.items),
+            'recordcount': int((instance+1) * workload_settings.inserts_per_workerinstance +
+                               workload_settings.items),
+        }
+
+    host = target.node
+    if target.cloud:
+        host = target.cloud.get("cluster_svc", host)
+
+    run_mongo_ycsb(
+        hosts=[host],
+        database=target.bucket,
+        username=target.username,
+        password=target.password,
+        action="run",
+        ycsb_client=workload_settings.ycsb_client,
+        workload=workload_settings.workload_path,
+        items=int(workload_settings.items),
+        workers=workload_settings.workers,
+        target=int(workload_settings.target),
+        ops=int(workload_settings.ops),
+        instance=instance,
+        execution_time=workload_settings.time,
+        timeseries=workload_settings.timeseries,
+        fieldlength=workload_settings.field_length,
+        fieldcount=workload_settings.field_count,
+        upsert=bool(workload_settings.upsert),
+        batch_size=workload_settings.mongo_batch_size or 1,
+        write_concern=workload_settings.mongo_durability_level or "majority",
+        requestdistribution=workload_settings.requestdistribution,
+        ycsb_jvm_args=workload_settings.ycsb_jvm_args,
+        insert_test_params=insert_test_params,
+        measurement_type=workload_settings.measurement_type,
+        histogram_buckets=workload_settings.histogram_buckets,
+        histogram_bucket_size=workload_settings.histogram_bucket_size,
+        verbose_histogram=bool(workload_settings.verbose_histogram),
+    )