apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
  name: ycsb
  namespace: argo
spec:
  ttlStrategy:
    secondsAfterCompletion: 120
  podGC:
    strategy: OnPodCompletion
  entrypoint: bench/adhoc
  serviceAccountName: workflow
  arguments:
    parameters:
    - name: suite
      value: "default"
    - name: tier
      value: t1.standard
    - name: region
      value: us-west-2
    - name: provider
      value: aws
    - name: repeats
      value: "1"
    - name: recordcount
      value: "10000000"
    - name: operationcount
      value: "1000000"
    - name: threadcount
      value: "100"
    - name: workload
      value: "a"
    - name: tikv-version
      value: "v4.0.7"
    - name: tidb-version
      value: "v4.0.7"
    - name: pd-version
      value: "v4.0.7"
    - name: tidb-config
      value: ""
    - name: pd-config
      value: ""
    - name: tikv-config
      value: ""

  templates:
  - name: bench/adhoc
    steps:
    - - name: setup
        templateRef:
          name: tidbcloud
          # with-cluster create a TiDB cluster which makes the workflow.outputs.parameters.tidb-XXXX available afterwards
          template: with-cluster
    - - name: load
        template: ycsb
        arguments:
          parameters:
          - name: action
            value: load
    - - name: run
        template: run-and-collect
        arguments:
          parameters:
          - name: ordinal
            value: '{{item}}'
        withSequence:
          count: '{{workflow.parameters.repeats}}'

  - name: run-and-collect
    inputs:
      parameters:
      - name: ordinal
        default: "1"
    steps:
    - - name: run
        template: ycsb
        arguments:
          paramters:
          - name: ordinal
            value: '{{inputs.parameters.ordinal}}'
          - name: action
            value: 'run'
    - - name: collect
        template: collect
        arguments:
          parameters:
          - name: mongo-secret
            value: "mongo-secret"
          artifacts:
          - name: result
            from: '{{steps.run.outputs.artifacts.job-log}}'

  - name: collect
    inputs:
      parameters:
      - name: mongo-secret
      artifacts:
      - name: result
        path: /tmp/result.log
    script:
      image: aylei/pymongo:latest
      command: [python]
      env:
      - name: MONGO_DSN
        valueFrom:
          secretKeyRef:
            name: '{{inputs.parameters.mongo-secret}}'
            key: dsn
      source: |
        import json
        import pymongo
        import re
        import os
        from datetime import datetime

        report = {}

        result = {}
        with open('/tmp/result.log') as f:
            content = f.read()
            sps = content.split("Run finished")
            if len(sps) < 2:
                sys.exit("cannot parse results")
            lines = sps[1].split("\n")
            for l in lines:
                m = re.match(r'([\s\S]+)\s+-\s(Takes.*)', l)
                if m is None:
                    continue
                o = {}
                pairs = m.group(2).split(',')
                for pair in pairs:
                    kv = pair.split(':')
                    o[kv[0].strip().replace('.', "")] = float(kv[1].strip())
                result[m.group(1)] = o
        report['result'] = result

        tidb = {
            'tidb': {
                'version': "{{workflow.parameters.tidb-version}}",
                'machine': "dummy"
            },
            'pd': {
                'version': "{{workflow.parameters.tidb-version}}",
                'machine': "dummy"
            },
            'tikv': {
                'version': "{{workflow.parameters.tidb-version}}",
                'machine': "dummy"
            },
            'tier': "{{workflow.parameters.tier}}",
            'provider': "{{workflow.parameters.provider}}",
            'region': "{{workflow.parameters.region}}",
            'id': "{{workflow.outputs.parameters.tidb-id}}",
        }
        # with open('tidb.json') as f:
        #     tc = json.load(f)
        #     spec = tc['spec']
        #     tidb['tidb']['config'] = spec['tidb']['config']
        #     tidb['tikv']['config'] = spec['tikv']['config']
        report['tidb'] = tidb

        ycsb = {}
        ycsb['recordcount'] = {{workflow.parameters.recordcount}}
        ycsb['threadcount'] = {{workflow.parameters.threadcount}}
        ycsb['workload'] = "{{workflow.parameters.workload}}"
        ycsb['operationcount'] = "{{workflow.parameters.operationcount}}"
        report['ycsb'] = ycsb

        report['time'] = datetime.now()
        report['suite'] = "{{workflow.parameters.suite}}"

        mc = pymongo.MongoClient(os.environ['MONGO_DSN'], ssl=True, ssl_cert_reqs='CERT_NONE')
        collection = mc["test"]["ycsb"]
        collection.insert_one(report)

  - name: ycsb
    inputs:
      parameters:
      - name: ordinal
        default: "1"
      - name: action
        default: "run"
    outputs:
      artifacts:
      - name: job-log
        from: '{{steps.execute.outputs.artifacts.job-log}}'
    steps:
    - - name: execute
        templateRef:
          name: utilities
          template: apply-job-and-wait
        arguments:
          parameters:
          - name: kubeconfig-secret
            value: '{{workflow.outputs.parameters.tidb-kubeconfig-secret}}'
          - name: timeout
            value: 6h
          - name: ns
            value: default
          artifacts:
          - name: manifest
            raw:
              data: |
                apiVersion: batch/v1
                kind: Job
                metadata:
                  name: ycsb{{inputs.parameters.ordinal}}-{{inputs.parameters.action}}-{{workflow.uid}}
                  labels:
                    app: ycsb
                spec:
                  backoffLimit: 0
                  parallelism: 1
                  template:
                    metadata:
                      labels:
                        app: ycsb
                    spec:
                      tolerations:
                      - effect: NoSchedule
                        key: dedicated
                        operator: Equal
                        value: benchmark
                      nodeSelector:
                        pool: benchmark
                      containers:
                        - name: ycsb
                          image: aylei/go-ycsb:20201029
                          command:
                            - /bin/sh
                            - -c
                            - |-
                              set -x
                              set -euo pipefail
                              /go-ycsb ${ACTION} mysql -P workloads/workload${WORKLOAD} \
                                -p mysql.host=${TIDB_HOST} \
                                -p mysql.port=${TIDB_PORT} \
                                -p mysql.password=${PASSWORD} \
                                -p recordcount=${RECORDCOUNT} \
                                -p operationcount=${OPERATIONCOUNT} \
                                -p threadcount=${THREADCOUNT}
                          env:
                            - name: TIDB_HOST
                              value: "db-tidb.tidb{{workflow.outputs.parameters.tidb-id}}"
                            - name: TIDB_PORT
                              value: "{{workflow.outputs.parameters.tidb-port}}"
                            - name: PASSWORD
                              value: "{{workflow.outputs.parameters.tidb-password}}"
                            - name: RECORDCOUNT
                              value: "{{workflow.parameters.recordcount}}"
                            - name: OPERATIONCOUNT
                              value: "{{workflow.parameters.operationcount}}"
                            - name: THREADCOUNT
                              value: "{{workflow.parameters.threadcount}}"
                            - name: WORKLOAD
                              value: "{{workflow.parameters.workload}}"
                            - name: ACTION
                              value: "{{inputs.parameters.action}}"
                      restartPolicy: Never