Rbroughan/pass logpath labels through (#9856)

tryangul · tryangul · commit 01b555e336b7 · 2023-11-15T10:04:47.000-08:00
diff --git a/airbyte-api/src/main/openapi/workload-openapi.yaml b/airbyte-api/src/main/openapi/workload-openapi.yaml
@@ -204,6 +204,7 @@ components:
       required:
       - id
       - inputPayload
+      - labels
       - logPath
       type: object
       properties:
@@ -222,7 +223,6 @@ components:
           nullable: true
         labels:
           type: array
-          nullable: true
           items:
             $ref: '#/components/schemas/WorkloadLabel'
         inputPayload:
@@ -254,6 +254,7 @@ components:
           type: string
     WorkloadCreateRequest:
       required:
+      - labels
       - logPath
       - workloadId
       - workloadInput
@@ -263,7 +264,6 @@ components:
           type: string
         labels:
           type: array
-          nullable: true
           items:
             $ref: '#/components/schemas/WorkloadLabel'
         workloadInput:
diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/WorkloadApiWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/WorkloadApiWorker.java
@@ -67,13 +67,14 @@ public ReplicationOutput run(final ReplicationInput replicationInput, final Path
 
     // TODO worker may resume, check if job exists first
     // Create the workload
-    createWorkload(new WorkloadCreateRequest(workloadId,
-        serializedInput,
-        jobRoot.toString(),
+    createWorkload(new WorkloadCreateRequest(
+        workloadId,
         List.of(
             new WorkloadLabel("connectionId", replicationInput.getConnectionId().toString()),
             new WorkloadLabel("jobId", replicationInput.getJobRunConfig().getJobId()),
-            new WorkloadLabel("attemptNumber", replicationInput.getJobRunConfig().getAttemptId().toString()))));
+            new WorkloadLabel("attemptNumber", replicationInput.getJobRunConfig().getAttemptId().toString())),
+        serializedInput,
+        jobRoot.toString()));
 
     // Wait until workload reaches a terminal status
     int i = 0;
diff --git a/airbyte-config/config-models/src/main/java/io/airbyte/config/messages/LauncherInputMessage.kt b/airbyte-config/config-models/src/main/java/io/airbyte/config/messages/LauncherInputMessage.kt
@@ -3,12 +3,32 @@ package io.airbyte.config.messages
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize
 
 @JsonDeserialize(builder = LauncherInputMessage.Builder::class)
-data class LauncherInputMessage(val workloadId: String, val workloadInput: String) {
-  data class Builder(var workloadId: String? = null, var workloadInput: String? = null) {
+data class LauncherInputMessage(
+  val workloadId: String,
+  val workloadInput: String,
+  val labels: Map<String, String>,
+  val logPath: String,
+) {
+  data class Builder(
+    var workloadId: String? = null,
+    var workloadInput: String? = null,
+    var labels: Map<String, String>? = null,
+    var logPath: String? = null,
+  ) {
     fun workloadId(workloadId: String) = apply { this.workloadId = workloadId }
 
     fun workloadInput(workloadInput: String) = apply { this.workloadInput = workloadInput }
 
-    fun build() = LauncherInputMessage(workloadId = workloadId!!, workloadInput = workloadInput!!)
+    fun labels(labels: Map<String, String>) = apply { this.labels = labels }
+
+    fun logPath(logPath: String) = apply { this.logPath = logPath }
+
+    fun build() =
+      LauncherInputMessage(
+        workloadId = workloadId!!,
+        workloadInput = workloadInput!!,
+        labels = labels!!,
+        logPath = logPath!!,
+      )
   }
 }
diff --git a/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/WorkloadApi.kt b/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/WorkloadApi.kt
@@ -79,7 +79,12 @@ open class WorkloadApi(
       workloadCreateRequest.workloadInput,
       workloadCreateRequest.logPath,
     )
-    workloadService.create(workloadId = workloadCreateRequest.workloadId, workloadInput = workloadCreateRequest.workloadInput)
+    workloadService.create(
+      workloadId = workloadCreateRequest.workloadId,
+      workloadInput = workloadCreateRequest.workloadInput,
+      workloadCreateRequest.labels.associate { it.key to it.value },
+      workloadCreateRequest.logPath,
+    )
   }
 
   @PUT
diff --git a/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/WorkloadService.kt b/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/WorkloadService.kt
@@ -28,12 +28,18 @@ open class WorkloadService(
   open fun create(
     workloadId: String,
     workloadInput: String,
+    labels: Map<String, String>,
+    logPath: String,
   ) {
     ApmTraceUtils.addTagsToTrace(mutableMapOf(WORKLOAD_ID_TAG to workloadId) as Map<String, Any>?)
-    messageProducer.publish("launcher-queue", LauncherInputMessage(workloadId, workloadInput))
+    messageProducer.publish(LAUNCHER_QUEUE_NAME, LauncherInputMessage(workloadId, workloadInput, labels, logPath))
     metricPublisher.count(
       WorkloadApiMetricMetadata.WORKLOAD_MESSAGE_PUBLISHED.metricName,
       MetricAttribute(WORKLOAD_ID_TAG, workloadId),
     )
   }
+
+  companion object {
+    const val LAUNCHER_QUEUE_NAME = "launcher-queue"
+  }
 }
diff --git a/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/domain/Workload.kt b/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/domain/Workload.kt
@@ -9,7 +9,7 @@ data class Workload(
   var dataplaneId: String? = null,
   var status: WorkloadStatus? = null,
   var lastHeartbeatAt: OffsetDateTime? = null,
-  var labels: MutableList<WorkloadLabel>? = null,
+  var labels: MutableList<WorkloadLabel> = mutableListOf(),
   var inputPayload: String = "",
   var logPath: String = "",
 )
diff --git a/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/domain/WorkloadCreateRequest.kt b/airbyte-workload-api-server/src/main/kotlin/io/airbyte/workload/api/domain/WorkloadCreateRequest.kt
@@ -5,7 +5,7 @@ import io.swagger.v3.oas.annotations.media.Schema
 data class WorkloadCreateRequest(
   @Schema(required = true)
   var workloadId: String = "",
-  var labels: List<WorkloadLabel>? = null,
+  var labels: List<WorkloadLabel> = listOf(),
   @Schema(required = true)
   var workloadInput: String = "",
   var logPath: String = "",
diff --git a/airbyte-workload-api-server/src/test/kotlin/io/airbyte/workload/api/WorkloadApiTest.kt b/airbyte-workload-api-server/src/test/kotlin/io/airbyte/workload/api/WorkloadApiTest.kt
@@ -70,10 +70,10 @@ class WorkloadApiTest(
   @Test
   fun `test create success`() {
     every { workloadHandler.createWorkload(any(), any(), any(), any()) } just Runs
-    every { workloadService.create(any(), any()) } just Runs
+    every { workloadService.create(any(), any(), any(), any()) } just Runs
     testEndpointStatus(HttpRequest.POST("/api/v1/workload/create", Jsons.serialize(WorkloadCreateRequest())), HttpStatus.NO_CONTENT)
     verify(exactly = 1) { workloadHandler.createWorkload(any(), any(), any(), any()) }
-    verify(exactly = 1) { workloadService.create(any(), any()) }
+    verify(exactly = 1) { workloadService.create(any(), any(), any(), any()) }
   }
 
   @Test
diff --git a/airbyte-workload-launcher/src/main/kotlin/StartupApplicationEventListener.kt b/airbyte-workload-launcher/src/main/kotlin/StartupApplicationEventListener.kt
@@ -9,23 +9,23 @@ import datadog.trace.api.Trace
 import io.airbyte.metrics.lib.ApmTraceUtils
 import io.airbyte.metrics.lib.MetricAttribute
 import io.airbyte.workload.api.client2.generated.WorkloadApi
-import io.airbyte.workload.api.client2.model.generated.Workload
 import io.airbyte.workload.api.client2.model.generated.WorkloadListRequest
 import io.airbyte.workload.api.client2.model.generated.WorkloadListResponse
 import io.airbyte.workload.api.client2.model.generated.WorkloadStatus
 import io.airbyte.workload.launcher.metrics.CustomMetricPublisher
 import io.airbyte.workload.launcher.metrics.MeterFilterFactory.Companion.DATA_PLANE_ID_TAG
-import io.airbyte.workload.launcher.metrics.MeterFilterFactory.Companion.REHYDRATION_OPERATION_NAME
+import io.airbyte.workload.launcher.metrics.MeterFilterFactory.Companion.RESUME_CLAIMED_OPERATION_NAME
 import io.airbyte.workload.launcher.metrics.MeterFilterFactory.Companion.WORKLOAD_ID_TAG
 import io.airbyte.workload.launcher.metrics.WorkloadLauncherMetricMetadata
+import io.airbyte.workload.launcher.model.toLauncherInput
 import io.airbyte.workload.launcher.pipeline.LaunchPipeline
-import io.airbyte.workload.launcher.pipeline.LauncherInput
 import io.github.oshai.kotlinlogging.KotlinLogging
 import io.micronaut.context.annotation.Value
 import io.micronaut.context.event.ApplicationEventListener
 import io.micronaut.discovery.event.ServiceReadyEvent
 import io.temporal.worker.WorkerFactory
 import jakarta.inject.Singleton
+import kotlin.concurrent.thread
 
 private val logger = KotlinLogging.logger {}
 
@@ -39,19 +39,20 @@ class StartupApplicationEventListener(
 ) :
   ApplicationEventListener<ServiceReadyEvent> {
   override fun onApplicationEvent(event: ServiceReadyEvent?) {
-    // TODO this might slowdown start quite a bit, should be reworked
-    try {
-      rehydrateAndProcessClaimed()
-    } catch (e: Exception) {
-      logger.error(e) { "rehydrateAndProcessClaimed failed" }
-    }
+    thread {
+      try {
+        retrieveAndProcessClaimed()
+      } catch (e: Exception) {
+        logger.error(e) { "rehydrateAndProcessClaimed failed" }
+      }
 
-    workerFactory.start()
+      workerFactory.start()
+    }
   }
 
   @VisibleForTesting
-  @Trace(operationName = REHYDRATION_OPERATION_NAME)
-  fun rehydrateAndProcessClaimed() {
+  @Trace(operationName = RESUME_CLAIMED_OPERATION_NAME)
+  fun retrieveAndProcessClaimed() {
     addTagsToTrace()
     val workloadListRequest =
       WorkloadListRequest(
@@ -63,20 +64,17 @@ class StartupApplicationEventListener(
       apiClient.workloadList(workloadListRequest)
 
     workloadList.workloads.forEach {
-      metricPublisher.count(WorkloadLauncherMetricMetadata.WORKLOAD_PROCESSED_ON_RESTART, MetricAttribute(WORKLOAD_ID_TAG, it.id))
-      pipe.accept(convertToInputMessage(it))
+      metricPublisher.count(
+        WorkloadLauncherMetricMetadata.WORKLOAD_PROCESSED_ON_RESTART,
+        MetricAttribute(WORKLOAD_ID_TAG, it.id),
+      )
+      pipe.accept(it.toLauncherInput())
     }
   }
 
   private fun addTagsToTrace() {
     val commonTags = hashMapOf<String, Any>()
-    commonTags.put(DATA_PLANE_ID_TAG, dataplaneId)
+    commonTags[DATA_PLANE_ID_TAG] = dataplaneId
     ApmTraceUtils.addTagsToTrace(commonTags)
   }
-
-  @VisibleForTesting
-  fun convertToInputMessage(workload: Workload): LauncherInput {
-    // TODO(Subodh): Add proper input once the format is decided
-    return LauncherInput(workload.id, "workload-input", "log-path.txt")
-  }
 }
diff --git a/airbyte-workload-launcher/src/main/kotlin/metrics/MeterFilterFactory.kt b/airbyte-workload-launcher/src/main/kotlin/metrics/MeterFilterFactory.kt
@@ -48,6 +48,6 @@ class MeterFilterFactory(
 
     const val LAUNCH_PIPELINE_OPERATION_NAME = "launch-pipeline"
     const val LAUNCH_PIPELINE_STAGE_OPERATION_NAME = "launch-pipeline-stage"
-    const val REHYDRATION_OPERATION_NAME = "rehydration"
+    const val RESUME_CLAIMED_OPERATION_NAME = "resume_claimed"
   }
 }
diff --git a/airbyte-workload-launcher/src/main/kotlin/metrics/WorkloadLauncherMetricMetadata.kt b/airbyte-workload-launcher/src/main/kotlin/metrics/WorkloadLauncherMetricMetadata.kt
@@ -16,6 +16,10 @@ enum class WorkloadLauncherMetricMetadata(
     "workload_received",
     "increments when the launcher receives a workload from the queue",
   ),
+  WORKLOAD_CLAIM_RESUMED(
+    "workload_claim_resumed",
+    "increments when a claimed workload is retrieved and processed on startup",
+  ),
   WORKLOAD_CLAIMED(
     "workload_claimed",
     "increments when the launcher claims a workload",
diff --git a/airbyte-workload-launcher/src/main/kotlin/mocks/MessageConsumer.kt b/airbyte-workload-launcher/src/main/kotlin/mocks/MessageConsumer.kt
@@ -36,7 +36,7 @@ class MessageConsumer(private val deserializer: PayloadDeserializer) {
     val id = mockJobId.getAndIncrement().toString()
     val mocked = mockJobIdFor(input, id)
 
-    return LauncherInput("workload-$id", mocked)
+    return LauncherInput("workload-$id", mocked, mapOf(), "/")
   }
 
   private fun mockJobIdFor(
diff --git a/airbyte-workload-launcher/src/main/kotlin/model/WorkloadExtensions.kt b/airbyte-workload-launcher/src/main/kotlin/model/WorkloadExtensions.kt
@@ -0,0 +1,13 @@
+package io.airbyte.workload.launcher.model
+
+import io.airbyte.workload.api.client2.model.generated.Workload
+import io.airbyte.workload.launcher.pipeline.LauncherInput
+
+fun Workload.toLauncherInput(): LauncherInput {
+  return LauncherInput(
+    this.id,
+    this.inputPayload,
+    this.labels.associate { it.key to it.value },
+    this.logPath,
+  )
+}
diff --git a/airbyte-workload-launcher/src/main/kotlin/pipeline/LaunchPipeline.kt b/airbyte-workload-launcher/src/main/kotlin/pipeline/LaunchPipeline.kt
@@ -76,7 +76,7 @@ class LaunchPipeline(
 
   private fun setLoggingScopeForWorkload(msg: LauncherInput): MdcScope {
     return MdcScope.Builder()
-      .setExtraMdcEntries(mapOf(LogClientSingleton.JOB_LOG_PATH_MDC_KEY to msg.jobLogPath))
+      .setExtraMdcEntries(mapOf(LogClientSingleton.JOB_LOG_PATH_MDC_KEY to msg.logPath))
       .build()
   }
 
diff --git a/airbyte-workload-launcher/src/main/kotlin/pipeline/LauncherInput.kt b/airbyte-workload-launcher/src/main/kotlin/pipeline/LauncherInput.kt
@@ -1,9 +1,8 @@
 package io.airbyte.workload.launcher.pipeline
 
-import java.io.File
-
 data class LauncherInput(
   val workloadId: String,
   val workloadInput: String,
-  val jobLogPath: String = File.createTempFile("log-path", ".txt").absolutePath,
+  val labels: Map<String, String>,
+  val logPath: String,
 )
diff --git a/airbyte-workload-launcher/src/main/kotlin/pipeline/LauncherMessageConsumer.kt b/airbyte-workload-launcher/src/main/kotlin/pipeline/LauncherMessageConsumer.kt
@@ -11,6 +11,8 @@ class LauncherMessageConsumer(private val launchPipeline: LaunchPipeline) : Mess
       LauncherInput(
         workloadId = input.workloadId,
         workloadInput = input.workloadInput,
+        labels = input.labels,
+        logPath = input.logPath,
       ),
     )
   }
diff --git a/airbyte-workload-launcher/src/test/kotlin/pipeline/LogPathTest.kt b/airbyte-workload-launcher/src/test/kotlin/pipeline/LogPathTest.kt
@@ -28,9 +28,9 @@ class LogPathTest {
     val dataplaneId = "dataplane_id"
 
     val logFile: File = File.createTempFile("log-path-1", ".txt")
-    val launcherInputMessage: LauncherInput =
-      LauncherInput(workloadId, "workload-input", logFile.absolutePath)
-    val launchStageIO: LaunchStageIO = LaunchStageIO(launcherInputMessage)
+    val launcherInputMessage =
+      LauncherInput(workloadId, "workload-input", mapOf(), logFile.absolutePath)
+    val launchStageIO = LaunchStageIO(launcherInputMessage)
 
     val claim: ClaimStage = claimStage(launchStageIO)
     val check: CheckStatusStage = checkStatusStage(launchStageIO)
@@ -39,7 +39,7 @@ class LogPathTest {
     val launch: LaunchPodStage = launchPodStage(launchStageIO, false)
     val statusUpdater: StatusUpdater = mockk()
 
-    val launchPipeline: LaunchPipeline =
+    val launchPipeline =
       LaunchPipeline(dataplaneId, claim, check, build, mutex, launch, statusUpdater, metricPublisher)
     launchPipeline.accept(launcherInputMessage)
 
@@ -52,9 +52,9 @@ class LogPathTest {
     val dataplaneId = "dataplane_id"
 
     val logFile: File = File.createTempFile("log-path-1", ".txt")
-    val launcherInputMessage: LauncherInput =
-      LauncherInput(workloadId, "workload-input", logFile.absolutePath)
-    val launchStageIO: LaunchStageIO = LaunchStageIO(launcherInputMessage)
+    val launcherInputMessage =
+      LauncherInput(workloadId, "workload-input", mapOf(), logFile.absolutePath)
+    val launchStageIO = LaunchStageIO(launcherInputMessage)
 
     val claim: ClaimStage = claimStage(launchStageIO)
     val check: CheckStatusStage = checkStatusStage(launchStageIO)
@@ -161,16 +161,16 @@ class LogPathTest {
     }
     val completeFileContent = fileContent.toString()
 
-    assert(completeFileContent.isNotBlank())
-    assert(completeFileContent.contains("CLAIM for workload 1"))
-    assert(completeFileContent.contains("CHECK_STATUS for workload 1"))
-    assert(completeFileContent.contains("BUILD for workload 1"))
-    assert(completeFileContent.contains("LAUNCH for workload 1"))
+    assert(completeFileContent.isNotBlank()) { "File content was blank" }
+    assert(completeFileContent.contains("CLAIM for workload 1")) { "CLAIM line missing" }
+    assert(completeFileContent.contains("CHECK_STATUS for workload 1")) { "CHECK_STATUS line missing" }
+    assert(completeFileContent.contains("BUILD for workload 1")) { "BUILD line missing" }
+    assert(completeFileContent.contains("LAUNCH for workload 1")) { "LAUNCH line missing" }
     if (assertException) {
-      assert(completeFileContent.contains("Failure for workload 1"))
-      assert(completeFileContent.contains("Completed without launching workload."))
+      assert(completeFileContent.contains("Failure for workload 1")) { "FAILURE line missing" }
+      assert(completeFileContent.contains("Completed without launching workload.")) { "TERMINAL line missing" }
     } else {
-      assert(completeFileContent.contains("Success: "))
+      assert(completeFileContent.contains("Success: ")) { "SUCCESS line missing" }
     }
   }
 }
diff --git a/airbyte-workload-launcher/src/test/kotlin/pipeline/PipelineStartupTest.kt b/airbyte-workload-launcher/src/test/kotlin/pipeline/PipelineStartupTest.kt
@@ -23,7 +23,9 @@ class PipelineStartupTest {
   fun `should process claimed workloads`() {
     val workloadId = "1"
     val dataplaneId = "US"
-    val launcherInput = LauncherInput(workloadId, "workload-input", "log-path.txt")
+    val payload = "payload"
+    val logPath = "/log/path"
+    val launcherInput = LauncherInput(workloadId, payload, mapOf(), logPath)
 
     val workloadApiClient: WorkloadApi = mockk()
     val workerFactory: WorkerFactory = mockk()
@@ -40,7 +42,7 @@ class PipelineStartupTest {
       )
 
     val workload =
-      Workload(workloadId, "payload", "/", dataplaneId, WorkloadStatus.cLAIMED)
+      Workload(workloadId, listOf(), payload, logPath, dataplaneId, WorkloadStatus.cLAIMED)
 
     val workloadListResponse = WorkloadListResponse(listOf(workload))
 
@@ -61,10 +63,9 @@ class PipelineStartupTest {
         ),
       )
 
-    startupApplicationEventListener.rehydrateAndProcessClaimed()
+    startupApplicationEventListener.retrieveAndProcessClaimed()
 
     verify { workloadApiClient.workloadList(workloadListRequest) }
-    verify { startupApplicationEventListener.convertToInputMessage(workload) }
     verify { launchPipeline.accept(launcherInput) }
   }
 }
diff --git a/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/BuildInputStageTest.kt b/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/BuildInputStageTest.kt
diff --git a/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/CheckStatusStageTest.kt b/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/CheckStatusStageTest.kt
diff --git a/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/ClaimStageTest.kt b/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/ClaimStageTest.kt
diff --git a/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/EnforceMutexStageTest.kt b/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/EnforceMutexStageTest.kt
diff --git a/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/LaunchPodStageTest.kt b/airbyte-workload-launcher/src/test/kotlin/pipeline/stages/LaunchPodStageTest.kt

Original file line number	Diff line number	Diff line change
`@@ -79,7 +79,12 @@ open class WorkloadApi(`
`79`	`79`	`workloadCreateRequest.workloadInput,`
`80`	`80`	`workloadCreateRequest.logPath,`
`81`	`81`	`)`
`82`		`- workloadService.create(workloadId = workloadCreateRequest.workloadId, workloadInput = workloadCreateRequest.workloadInput)`
	`82`	`+ workloadService.create(`
	`83`	`+ workloadId = workloadCreateRequest.workloadId,`
	`84`	`+ workloadInput = workloadCreateRequest.workloadInput,`
	`85`	`+ workloadCreateRequest.labels.associate { it.key to it.value },`
	`86`	`+ workloadCreateRequest.logPath,`
	`87`	`+ )`
`83`	`88`	`}`
`84`	`89`
`85`	`90`	`@PUT`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ data class Workload(`
`9`	`9`	`var dataplaneId: String? = null,`
`10`	`10`	`var status: WorkloadStatus? = null,`
`11`	`11`	`var lastHeartbeatAt: OffsetDateTime? = null,`
`12`		`- var labels: MutableList<WorkloadLabel>? = null,`
	`12`	`+ var labels: MutableList<WorkloadLabel> = mutableListOf(),`
`13`	`13`	`var inputPayload: String = "",`
`14`	`14`	`var logPath: String = "",`
`15`	`15`	`)`
Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,6 @@ class MeterFilterFactory(`
`48`	`48`
`49`	`49`	`const val LAUNCH_PIPELINE_OPERATION_NAME = "launch-pipeline"`
`50`	`50`	`const val LAUNCH_PIPELINE_STAGE_OPERATION_NAME = "launch-pipeline-stage"`
`51`		`- const val REHYDRATION_OPERATION_NAME = "rehydration"`
	`51`	`+ const val RESUME_CLAIMED_OPERATION_NAME = "resume_claimed"`
`52`	`52`	`}`
`53`	`53`	`}`