import groovy.json.JsonOutput

def main() {
    def clusterName = env.NODE_NAME
    def clusterValues = """#@data/values
---
name: 'jepsen'
"""
    def clusterNodes = """#@ load("@ytt:data", "data")
---
apiVersion: v1
kind: Service
metadata:
  name: #@ data.values.name
  labels:
    app: jenkins-jepsen-test
spec:
  type: ClusterIP
  clusterIP: None
  selector:
    name: #@ data.values.name
  ports:
  - name: ssh
    port: 22
    protocol: TCP
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: #@ data.values.name
  labels:
    app: jenkins-jepsen-test
spec:
  replicas: 5
  serviceName: #@ data.values.name
  selector:
    matchLabels:
      name: #@ data.values.name
  template:
    metadata:
      labels:
        name: #@ data.values.name
    spec:
      containers:
      - name: node
        image: hub.pingcap.net/qa/jepsen-node
        securityContext:
          privileged: true
        env:
        - name: AUTHORIZED_KEYS
          value: #@ data.read('key.pub')
        resources:
          requests:
            cpu: '2000m'
            memory: '4Gi'
          limits:
            cpu: '8000m'
            memory: '16Gi'
"""

    def peers = [0, 1, 2, 3, 4]
    try {
        stage("Setup") {
            container('control') {
                sh("chown -R 1000:1000 ./")
                deleteDir()
                sh("""
                mkdir -p manifests resources
                chown -R 1000:1000 manifests resources
                ssh-keygen -b 2048 -t rsa -m pem -f manifests/key -q -N ''
                """)
                dir("manifests") {
                    writeFile(file: "values.yaml", text: clusterValues)
                    writeFile(file: "nodes.yaml", text: clusterNodes)

                }
                dir("resources") {
                    if (params.TIDB_CONFIG) { writeFile(file: "tidb.conf", text: params.TIDB_CONFIG) }
                    if (params.TIKV_CONFIG) { writeFile(file: "tikv.conf", text: params.TIKV_CONFIG) }
                    if (params.PD_CONFIG) { writeFile(file: "pd.conf", text: params.PD_CONFIG) }
                }
            }
            container('kubectl') {
                sh("ytt -f manifests -v name=$clusterName | tee cluster.yaml")
                sh("kubectl apply -f cluster.yaml")
            }
            container('control') {
                timeout(10) {
                    peers.each {
                        sh("while ! ssh -i manifests/key -o StrictHostKeyChecking=no $clusterName-$it.$clusterName date ; do sleep 1; done")
                    }
                }
            }
        }
        stage("Test") {
            container('control') {
                def cmd = "java -cp resources:/jepsen-tidb.jar tidb.core test --os debian --username=root --ssh-private-key=manifests/key"
                cmd += " --nodes=${peers.collect { "$clusterName-$it.$clusterName" } .join(",")}"
                if (params.RAW_ARGS) {
                    cmd += " " + params.RAW_ARGS
                } else {
                    cmd += " --version='${params.TITLE}'"
                    cmd += " --tarball-url='${params.TARBALL_URL}'"
                    cmd += " --time-limit=${params.TIME_LIMIT}"
                    cmd += " --concurrency=${params.CONCURRENCY}"
                    cmd += " --nemesis=${params.NEMESIS}"
                    cmd += " --nemesis-interval=${params.NEMESIS_INTERVAL}"
                    cmd += " --workload=${params.WORKLOAD}"
                    cmd += " --txn-mode=${params.TXN_MODE}"
                    cmd += " --isolation=${params.ISOLATION}"
                    if (params.READ_LOCK) { cmd += " --read-lock=${params.READ_LOCK}" }
                    if (params.USE_INDEX) { cmd += " --use-index=${params.USE_INDEX}" }
                    if (params.UPDATE_IN_PLACE) { cmd += " --update-in-place=${params.UPDATE_IN_PLACE}" }
                    if (params.FORCE_REINSTALL) { cmd += " --force-reinstall=${params.FORCE_REINSTALL}" }
                    if (params.INIT_SQL) { cmd += " --init-sql=\"${params.INIT_SQL}\""}
                    if (params.INIT_TXN_SQL) { cmd += " --init-txn-sql=\"${params.INIT_TXN_SQL}\""}
                    if (params.EXTRA_ARGS) { cmd += " " + params.EXTRA_ARGS }
                }

                def result = [
                    id: UUID.randomUUID().toString(),
                    name: "jepsen-tidb-${params.TITLE}",
                    started_at: System.currentTimeMillis()/1000,
                    annotations: [
                        "jepsen.image": params.IMAGE,
                        "jenkins.build": BUILD_URL,
                        "jenkins.started_at": new Date(currentBuild.startTimeInMillis).format("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'"),
                    ],
                    meta: [
                        components: ["tidb/transaction"],
                        maintainers: ["qa@pingcap.com"],
                    ],
                    params: params,
                ]

                def resultV2 = [
                    name: "jepsen-tidb-${params.TITLE}",
                    owner: "guanzhongyang@pingcap.com",
                    started_at: (System.currentTimeMillis()/1000).intValue(),
                    external_id: BUILD_TAG,
                    details_url: BUILD_URL,
                    labels: [
                        "components": "transaction",
                        "jepsen.image": params.IMAGE.toString(),
                        "jepsen.title": params.TITLE.toString(),
                        "jepsen.params.tarball-url": params.TARBALL_URL.toString(),
                        "jepsen.params.time-limit": params.TIME_LIMIT.toString(),
                        "jepsen.params.concurrency": params.CONCURRENCY.toString(),
                        "jepsen.params.nemesis": params.NEMESIS.toString(),
                        "jepsen.params.nemesis-interval": params.NEMESIS_INTERVAL.toString(),
                        "jepsen.params.workload": params.WORKLOAD.toString(),
                        "jepsen.params.txn-mode": params.TXN_MODE.toString(),
                        "jepsen.params.isolation": params.ISOLATION.toString(),
                        "jepsen.params.read-lock": params.READ_LOCK.toString(),
                        "jepsen.params.use-index": params.USE_INDEX.toString(),
                        "jepsen.params.update-in-place": params.UPDATE_IN_PLACE.toString(),
                        "jepsen.params.force-reinstall": params.FORCE_REINSTALL.toString(),
                        "jepsen.params.init-sql": params.INIT_SQL.toString(),
                        "jepsen.params.init-txn-sql": params.INIT_TXN_SQL.toString(),
                    ]
                ]

                try {
                    def lastErr = null
                    for (trials = 0; trials < 5; trials++) {
                        try {
                            try {
                                sh("pkill -9 tcpdump || mkdir -p store")
                                peers.each {
                                    sh("tcpdump -i any -w store/node-${it}.pcap tcp port 4000 and host ${clusterName}-${it}.${clusterName} &")
                                }
                            } catch (ee) { echo "${ee}" }
                            sh("timeout ${params.TIME_LIMIT.toInteger() * 2 + 300} ${cmd}")
                            result.annotations['jepsen.trials'] = trials
                            resultV2.labels['jepsen.trials'] = trials.toString()
                            lastErr = null
                            break
                        } catch (e) {
                            try { sh("pkill tcpdump") } catch (ee) { echo "${ee}" }
                            if (e.toString().contains("exit code 124") || e.toString().contains("exit code 255")) {
                                lastErr = e
                                sleep(time: trials*2+1, unit: 'SECONDS')
                                continue
                            }
                            throw e
                        }
                    }
                    if (lastErr) { throw lastErr }
                    result.finished_at = System.currentTimeMillis()/1000
                    resultV2.completed_at = (System.currentTimeMillis()/1000).intValue()
                    resultV2.status = "completed"
                    result.status = "PASSED"
                    resultV2.conclusion = "success"
                    if (params.ALWAYS_ARCHIVE_STORE) { archiveArtifacts(artifacts: "store/**/*") }
                } catch (e) {
                    result.finished_at = System.currentTimeMillis()/1000
                    resultV2.completed_at = (System.currentTimeMillis()/1000).intValue()
                    resultV2.status = "completed"
                    result.reason = e.toString()
                    resultV2.output = e.toString()
                    if (e instanceof org.jenkinsci.plugins.workflow.steps.FlowInterruptedException) {
                        // may be timeout or canceled
                        result.status = "SKIPPED"
                        resultV2.conclusion = "skipped"
                    } else {
                        result.status = "FAILED"
                        resultV2.conclusion = "failure"
                    }
                    // TODO: alert if exit with 1
                    archiveArtifacts(artifacts: "store/**/*")
                    throw e
                } finally {
                    echo(JsonOutput.prettyPrint(JsonOutput.toJson(resultV2)))
                    if (params.REPORT) {
                        withCredentials([usernameColonPassword(credentialsId: "result-store", variable: "AUTH")]) {
                            httpRequest(url: "http://$AUTH@${env.RESULT_STORE_ENDPOINT}${result.id}",
                                        httpMode: "PUT", contentType: "APPLICATION_JSON",
                                        requestBody: JsonOutput.toJson(result))
                        }
                        // httpRequest(url: "https://testing.pingcap.net/records",
                        //             httpMode: "POST", contentType: "APPLICATION_JSON",
                        //             requestBody: JsonOutput.toJson(resultV2))
                    }
                }
            }
        }
    } finally {
        stage("Teardown") {
            container('kubectl') {
                if (fileExists('cluster.yaml')) {
                    sh("kubectl delete -f cluster.yaml")
                }
            }
        }
    }
}

def run(label, image, Closure body) {
    podTemplate(name: label, label: label, instanceCap: 10, serviceAccount: 'jepsen-runner', containers: [
        containerTemplate(name: 'control', image: image, alwaysPullImage: true, ttyEnabled: true, command: 'cat'),
        containerTemplate(name: 'kubectl', image: 'hub.pingcap.net/qa/kubetools:20200730', alwaysPullImage: false, ttyEnabled: true, command: 'cat'),
    ]) { node(label) { body() } }
}

catchError {
    run("jepsen-test-${BUILD_NUMBER}", params.IMAGE) { main() }
}
