apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
  name: tpcc
  namespace: argo
  labels:
    type: bench
    benchmark: tpcc
spec:
  ttlStrategy:
    secondsAfterCompletion: 120
  podGC:
    strategy: OnPodCompletion
  entrypoint: bench/adhoc
  serviceAccountName: workflow
  arguments:
    parameters:
    - name: tier
      value: t1.standard
    - name: region
      value: us-west-2
    - name: provider
      value: aws
    - name: repeats
      value: "1"
    - name: warehouses
      value: "1000"
    - name: duration
      value: 10m
    - name: terminals
      value: "100"
    - 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: ""
    - name: suite
      value: "default"

  templates:
  - name: bench/10K
    steps:
    - - name: setup
        templateRef:
          name: tidbcloud
          template: with-cluster
    - - name: restore
        templateRef:
          name: dataset
          template: tpcc-10000
        arguments:
          parameters:
          - name: kubeconfig-secret
            value: '{{workflow.outputs.parameters.tidb-kubeconfig-secret}}'
          - name: ns
            value: 'tidb{{workflow.outputs.parameters.tidb-id}}'
    - - name: bench
        template: bench-cycle
        arguments:
          parameters:
          - name: repeats
            value: '{{workflow.parameters.repeats}}'
          - name: warehouses
            value: '{{steps.restore.outputs.parameters.scale}}'

  - 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: bench
        template: adhoc-cycle
        arguments:
          parameters:
          - name: repeats
            value: '{{workflow.parameters.repeats}}'

  - name: adhoc-cycle
    inputs:
      parameters:
      - name: repeats
        default: "1"
    parallelism: 1
    steps:
    - - name: prepare
        template: tpcc
        arguments:
          parameters:
          - name: action
            value: 'prepare'
    - - name: loop
        template: run-and-collect
        arguments:
          parameters:
          - name: id
            value: "{{item}}"
        withSequence:
          count: '{{inputs.parameters.repeats}}'
    - - name: check
        template: tpcc
        arguments:
          parameters:
          - name: action
            value: 'check'

  - name: bench-cycle
    inputs:
      parameters:
      - name: repeats
        default: "1"
      - name: warehouse
        default: '{{workflow.parameters.warehouses}}'
    parallelism: 1
    steps:
    - - name: loop
        template: run-and-collect
        arguments:
          parameters:
          - name: id
            value: "{{item}}"
          - name: warehouses
            value: '{{inputs.parameters.warehouse}}'
        withSequence:
          count: '{{inputs.parameters.repeats}}'
    - - name: check
        template: tpcc
        arguments:
          parameters:
          - name: action
            value: 'check'
          - name: warehouses
            value: '{{inputs.parameters.warehouse}}'

  - name: run-and-collect
    inputs:
      parameters:
      - name: id
        default: "0"
      - name: warehouses
        default: '{{workflow.parameters.warehouses}}'
    steps:
    - - name: run
        template: tpcc
        arguments:
          parameters:
          - name: action
            value: 'run'
          - name: id
            value: '{{inputs.parameters.id}}'
          - name: warehouses
            value: '{{inputs.parameters.warehouses}}'
    - - name: save
        template: collect
        arguments:
          parameters:
          - name: mongo-secret
            value: "mongo-secret"
          - name: warehouses
            value: '{{inputs.parameters.warehouses}}'
          artifacts:
          - name: result
            from: '{{steps.run.outputs.artifacts.job-log}}'

  - name: tpcc
    inputs:
      parameters:
        - name: action
          default: run
        - name: id
          default: ""
    steps:
    - - name: generate-job
        template: gen-job
        arguments:
          parameters:
          - name: action
            value: '{{inputs.parameters.action}}'
          - name: id
            value: '{{inputs.parameters.id}}'
    - - name: run-job
        templateRef:
          name: utilities
          template: apply-job-and-wait
        arguments:
          parameters:
          - name: kubeconfig-secret
            value: '{{workflow.outputs.parameters.tidb-kubeconfig-secret}}'
          - name: timeout
            value: 12h
          - name: ns
            value: default
          artifacts:
          - name: manifest
            from: '{{steps.generate-job.outputs.artifacts.job-spec}}'
    outputs:
      artifacts:
      - name: job-log
        from: '{{steps.run-job.outputs.artifacts.job-log}}'

  - name: collect
    inputs:
      parameters:
      - name: mongo-secret
      - name: warehouses
      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()
            lines = content.split("\n")
            for l in lines:
                m = re.match(r'\[Summary\]\s([\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

        tpcc = {}
        tpcc['warehouses'] = {{inputs.parameters.warehouses}}
        tpcc['terminals'] = {{workflow.parameters.terminals}}
        tpcc['duration'] = "{{workflow.parameters.duration}}"
        report['tpcc'] = tpcc

        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"]["tpcc"]
        collection.insert_one(report)


  - name: gen-job
    inputs:
      parameters:
      - name: action
        default: run
      - name: id
        default: ''
      - name: warehouses
        default: '{{workflow.parameters.warehouses}}'
    outputs:
      artifacts:
      - name: job-spec
        path: /tmp/job.yaml
    script:
      image: bash
      command: [bash]
      # TODO: set owner reference (in shoot) for jobs to get them cleaned with the workflow
      source: |
        cat << EOF > /tmp/job.yaml
        apiVersion: batch/v1
        kind: Job
        metadata:
          name: tpcc{{inputs.parameters.id}}-{{inputs.parameters.action}}-{{workflow.uid}}
          labels:
            app: tpcc
        spec:
          backoffLimit: 0
          parallelism: 1
          template:
            metadata:
              labels:
                app: tpcc
            spec:
              nodeSelector:
                pool: benchmark
              tolerations:
              - effect: NoSchedule
                key: dedicated
                operator: Equal
                value: benchmark
              containers:
                - name: tpcc
                  image: aylei/go-tpc:20201027
                  command:
                    - /bin/sh
                    - -c
                    - |-
                      set -x
                      set -euo pipefail
                      elapseTime=0
                      threshold=600
                      while true; do
                          sleep 5
                          elapseTime=\$(( elapseTime+5 ))

                          if [[ \${elapseTime} -ge \${threshold} ]]
                          then
                              echo "waiting for TiDB endpoint ready timeout" >&2
                              exit 1
                          fi

                          if nslookup \${TIDB_HOST} 2>/dev/null
                          then
                              echo "nslookup domain \${TIDB_HOST} success"
                              break
                          else
                              echo "nslookup domain \${TIDB_HOST} failed" >&2
                          fi
                      done
                      case \${ACTION} in
                        prepare)
                          mysql -h\${TIDB_HOST} -P\${TIDB_PORT} -p\${PASSWORD} -e "DROP DATABASE IF EXISTS tpcc; CREATE DATABASE tpcc;"
                          /go-tpc -H=\${TIDB_HOST} -P=\${TIDB_PORT} -p=\${PASSWORD} -D=tpcc --warehouses=\${WAREHOUSES} tpcc prepare
                          ;;
                        run)
                          /go-tpc -H=\${TIDB_HOST} -P=\${TIDB_PORT} -p=\${PASSWORD} -D=tpcc --warehouses=\${WAREHOUSES} --time=\${DURATION} --ignore-error tpcc run
                          ;;
                        check)
                          /go-tpc -H=\${TIDB_HOST} -P=\${TIDB_PORT} -p=\${PASSWORD} -D=tpcc --warehouses=\${WAREHOUSES} tpcc check
                          ;;
                      esac
                  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: WAREHOUSES
                      value: "{{inputs.parameters.warehouses}}"
                    - name: DURATION
                      value: "{{workflow.parameters.duration}}"
                    - name: TERMINALS
                      value: "{{workflow.parameters.terminals}}"
                    - name: ACTION
                      value: "{{inputs.parameters.action}}"
              restartPolicy: Never
        EOF