Merge pull request #65 from leanix/feature/CID-1287/heartbeat

CID:1287 Adds heartbeat and graceful shutdown

CODEOWNERS

@@ -1,4 +1,4 @@
 # These owners will be the default owners for everything in
 # the repo. Unless a later match takes precedence,
 # they will be requested for review when someone opens a pull request.
-* @alfredo-mfaria @aravindmetku @geoandri @henriqamaral @mohamedlajmileanix
+* @leanix/team-cider

build.gradle.kts

@@ -1,13 +1,13 @@
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 
 plugins {
-    id("org.springframework.boot") version "3.0.5"
+    id("org.springframework.boot") version "3.1.0"
     id("io.spring.dependency-management") version "1.1.0"
-    id("io.gitlab.arturbosch.detekt") version "1.21.0"
+    id("io.gitlab.arturbosch.detekt") version "1.23.0"
     id("com.expediagroup.graphql") version "6.3.1"
     id("org.cyclonedx.bom") version "1.7.2"
-    kotlin("jvm") version "1.6.21"
-    kotlin("plugin.spring") version "1.6.21"
+    kotlin("jvm") version "1.8.21"
+    kotlin("plugin.spring") version "1.8.21"
     jacoco
 }
 
@@ -45,7 +45,7 @@ dependencies {
 
 dependencyManagement {
     imports {
-        mavenBom("org.springframework.cloud:spring-cloud-dependencies:2022.0.1")
+        mavenBom("org.springframework.cloud:spring-cloud-dependencies:2022.0.3")
     }
     dependencies {
         dependency("com.google.guava:guava:30.0-jre")

src/main/kotlin/net/leanix/vsm/githubbroker/connector/adapter/feign/FeignRunStatusProvider.kt

@@ -0,0 +1,13 @@
+package net.leanix.vsm.githubbroker.connector.adapter.feign
+
+import net.leanix.vsm.githubbroker.connector.adapter.feign.data.UpdateRunStateRequest
+import net.leanix.vsm.githubbroker.connector.domain.RunStatusProvider
+import org.springframework.stereotype.Component
+
+@Component
+class FeignRunStatusProvider(private val vsmClient: VsmClient) : RunStatusProvider {
+
+    override fun updateRunStatus(runId: String, runState: UpdateRunStateRequest) {
+        vsmClient.updateRunState(runId, runState)
+    }
+}

src/main/kotlin/net/leanix/vsm/githubbroker/connector/adapter/feign/VsmClient.kt

@@ -6,11 +6,13 @@ import net.leanix.vsm.githubbroker.connector.adapter.feign.data.DoraRequest
 import net.leanix.vsm.githubbroker.connector.adapter.feign.data.LanguageRequest
 import net.leanix.vsm.githubbroker.connector.adapter.feign.data.ServiceRequest
 import net.leanix.vsm.githubbroker.connector.adapter.feign.data.TopicRequest
+import net.leanix.vsm.githubbroker.connector.adapter.feign.data.UpdateRunStateRequest
 import net.leanix.vsm.githubbroker.shared.Constants.EVENT_TYPE_HEADER
 import net.leanix.vsm.githubbroker.shared.auth.adapter.feign.config.MtmFeignClientConfiguration
 import org.springframework.cloud.openfeign.FeignClient
 import org.springframework.web.bind.annotation.GetMapping
 import org.springframework.web.bind.annotation.PostMapping
+import org.springframework.web.bind.annotation.PutMapping
 import org.springframework.web.bind.annotation.RequestBody
 import org.springframework.web.bind.annotation.RequestHeader
 import org.springframework.web.bind.annotation.RequestParam
@@ -47,4 +49,13 @@ interface VsmClient {
         @RequestParam("integrationName") integrationName: String,
         @RequestParam("configSetName") configSetName: String
     ): AssignmentResponse
+
+    @GetMapping("/health/heartbeat")
+    fun heartbeat(@RequestParam("runId") runId: String): String
+
+    @PutMapping("/run/status")
+    fun updateRunState(
+        @RequestParam("runId") runId: String,
+        @RequestBody runState: UpdateRunStateRequest,
+    )
 }

src/main/kotlin/net/leanix/vsm/githubbroker/connector/adapter/feign/data/UpdateRunStateRequest.kt

@@ -0,0 +1,23 @@
+package net.leanix.vsm.githubbroker.connector.adapter.feign.data
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+data class UpdateRunStateRequest(
+    val state: RunState,
+    val workspaceId: String? = null,
+    val connector: String? = null,
+    val orgName: String? = null,
+    val message: String? = null,
+    val region: String? = null
+)
+
+enum class RunState {
+    QUEUED,
+    DELETED,
+    RUNNING,
+    FINISHED,
+    FAILED,
+    FINISHED_FOR_LIVE,
+    LIVE
+}

src/main/kotlin/net/leanix/vsm/githubbroker/connector/domain/RunStatusProvider.kt

@@ -0,0 +1,7 @@
+package net.leanix.vsm.githubbroker.connector.domain
+
+import net.leanix.vsm.githubbroker.connector.adapter.feign.data.UpdateRunStateRequest
+
+interface RunStatusProvider {
+    fun updateRunStatus(runId: String, runState: UpdateRunStateRequest)
+}

src/main/kotlin/net/leanix/vsm/githubbroker/connector/runner/InitialStateRunner.kt

@@ -9,6 +9,7 @@ import net.leanix.vsm.githubbroker.connector.domain.CommandProvider
 import net.leanix.vsm.githubbroker.logs.application.LoggingService
 import net.leanix.vsm.githubbroker.logs.domain.LogStatus
 import net.leanix.vsm.githubbroker.logs.domain.StatusLog
+import net.leanix.vsm.githubbroker.shared.cache.AssignmentCache
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 import org.springframework.boot.ApplicationArguments
@@ -54,7 +55,9 @@ class InitialStateRunner(
     }
     private fun getAssignments(): List<Assignment>? {
         kotlin.runCatching {
-            return assignmentService.getAssignments()
+            val assignments = assignmentService.getAssignments()
+            AssignmentCache.addAll(assignments)
+            return assignments
         }.onFailure {
             logger.error("Failed to get initial state. No assignment found for this workspace id")
         }

src/main/kotlin/net/leanix/vsm/githubbroker/connector/runner/ShutdownService.kt

@@ -0,0 +1,36 @@
+package net.leanix.vsm.githubbroker.connector.runner
+
+import jakarta.annotation.PreDestroy
+import net.leanix.vsm.githubbroker.connector.adapter.feign.FeignRunStatusProvider
+import net.leanix.vsm.githubbroker.connector.adapter.feign.data.RunState
+import net.leanix.vsm.githubbroker.connector.adapter.feign.data.UpdateRunStateRequest
+import net.leanix.vsm.githubbroker.shared.Constants.GITHUB_ENTERPRISE_CONNECTOR
+import net.leanix.vsm.githubbroker.shared.cache.AssignmentCache
+import org.slf4j.LoggerFactory
+import org.springframework.stereotype.Service
+
+@Service
+class ShutdownService(private val runStatusProvider: FeignRunStatusProvider) {
+
+    private val logger = LoggerFactory.getLogger(ShutdownService::class.java)
+
+    @PreDestroy
+    fun onDestroy() {
+        if (AssignmentCache.getAll().isEmpty()) {
+            logger.info("Shutting down github broker before receiving any assignment")
+        } else {
+            AssignmentCache.getAll().values.forEach { assignment ->
+                runStatusProvider.updateRunStatus(
+                    assignment.runId.toString(),
+                    UpdateRunStateRequest(
+                        state = RunState.FINISHED,
+                        workspaceId = assignment.workspaceId.toString(),
+                        connector = GITHUB_ENTERPRISE_CONNECTOR,
+                        orgName = assignment.organizationName,
+                        message = "Gracefully stopped github enterprise"
+                    )
+                )
+            }
+        }
+    }
+}

src/main/kotlin/net/leanix/vsm/githubbroker/connector/scheduler/HeartbeatScheduler.kt

@@ -0,0 +1,17 @@
+package net.leanix.vsm.githubbroker.connector.scheduler
+
+import net.leanix.vsm.githubbroker.connector.adapter.feign.VsmClient
+import net.leanix.vsm.githubbroker.shared.cache.AssignmentCache
+import org.springframework.scheduling.annotation.Scheduled
+import org.springframework.stereotype.Component
+
+@Component
+class HeartbeatScheduler(
+    private val vsmClient: VsmClient
+) {
+
+    @Scheduled(fixedRate = 300000) // 5 minutes
+    fun heartbeat() {
+        AssignmentCache.getAll().values.forEach { assigment -> vsmClient.heartbeat(assigment.runId.toString()) }
+    }
+}

src/main/kotlin/net/leanix/vsm/githubbroker/connector/scheduler/ScheduleRepositories.kt

@@ -5,6 +5,7 @@ import net.leanix.vsm.githubbroker.connector.application.RepositoriesService
 import net.leanix.vsm.githubbroker.connector.domain.Assignment
 import net.leanix.vsm.githubbroker.connector.domain.CommandEventAction
 import net.leanix.vsm.githubbroker.connector.domain.CommandProvider
+import net.leanix.vsm.githubbroker.shared.cache.AssignmentCache
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 import org.springframework.scheduling.annotation.Scheduled
@@ -34,7 +35,9 @@ class ScheduleRepositories(
 
     private fun getAssignments(): List<Assignment>? {
         return kotlin.runCatching {
-            return assignmentService.getAssignments()
+            val assignments = assignmentService.getAssignments()
+            AssignmentCache.addAll(assignments)
+            return assignments
         }.getOrElse {
             logger.error("Failed to get initial state. No assignment found for this workspace id")
             null

src/main/kotlin/net/leanix/vsm/githubbroker/shared/cache/AssignmentCache.kt

@@ -0,0 +1,24 @@
+package net.leanix.vsm.githubbroker.shared.cache
+
+import net.leanix.vsm.githubbroker.connector.domain.Assignment
+
+object AssignmentCache {
+
+    private val assigmentCache: MutableMap<String, Assignment> = mutableMapOf()
+
+    fun addAll(newAssignments: List<Assignment>) {
+        newAssignments.forEach { assignment -> assigmentCache[assignment.organizationName] = assignment }
+    }
+
+    fun get(key: String): Assignment? {
+        return assigmentCache[key]
+    }
+
+    fun getAll(): Map<String, Assignment> {
+        return assigmentCache
+    }
+
+    fun deleteAll() {
+        assigmentCache.clear()
+    }
+}

src/test/kotlin/net/leanix/vsm/githubbroker/connector/runner/ShutdownServiceTest.kt

@@ -0,0 +1,60 @@
+package net.leanix.vsm.githubbroker.connector.runner
+
+import io.mockk.clearAllMocks
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.slot
+import io.mockk.verify
+import net.leanix.vsm.githubbroker.connector.adapter.feign.FeignRunStatusProvider
+import net.leanix.vsm.githubbroker.connector.adapter.feign.data.RunState.FINISHED
+import net.leanix.vsm.githubbroker.connector.adapter.feign.data.UpdateRunStateRequest
+import net.leanix.vsm.githubbroker.connector.domain.Assignment
+import net.leanix.vsm.githubbroker.shared.Constants.GITHUB_ENTERPRISE_CONNECTOR
+import net.leanix.vsm.githubbroker.shared.cache.AssignmentCache
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import java.util.*
+
+class ShutdownServiceTest {
+    private val runStatusProvider: FeignRunStatusProvider = mockk(relaxed = true)
+    private val shutdownService = ShutdownService(runStatusProvider)
+
+    @BeforeEach
+    fun setup() {
+        clearAllMocks()
+        AssignmentCache.deleteAll()
+    }
+
+    @Test
+    fun `test onDestroy with empty assignment cache`() {
+        shutdownService.onDestroy()
+
+        // Assert no interactions with the run status provider
+        verify(exactly = 0) { runStatusProvider.updateRunStatus(any(), any()) }
+    }
+
+    @Test
+    fun `test onDestroy with assignments`() {
+        val assignment = Assignment(UUID.randomUUID(), UUID.randomUUID(), UUID.randomUUID(), null, "mock-org")
+        AssignmentCache.addAll(listOf(assignment))
+
+        val runStateSlot = slot<UpdateRunStateRequest>()
+
+        every { runStatusProvider.updateRunStatus(any(), capture(runStateSlot)) } answers { }
+
+        shutdownService.onDestroy()
+
+        verify {
+            runStatusProvider.updateRunStatus(
+                eq(assignment.runId.toString()),
+                match {
+                    it.state == FINISHED &&
+                        it.workspaceId == assignment.workspaceId.toString() &&
+                        it.connector == GITHUB_ENTERPRISE_CONNECTOR &&
+                        it.orgName == assignment.organizationName &&
+                        it.message == "Gracefully stopped github enterprise"
+                }
+            )
+        }
+    }
+}