From 6fb3e6d1a967644ad6de5ecad16719c72f7f34f0 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Mon, 19 Oct 2026 15:51:54 +0200 Subject: [PATCH] Schedule background tasks before sending start event --- .../background/BackgroundProcess.java | 15 +++++++-------- .../background/HeartbeatSchedulingTest.java | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java b/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java index f3ddbc8c..614a07b8 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java @@ -45,14 +45,6 @@ public void run() { ReportingApiHTTP api = new ReportingApiHTTP(getAikidoAPIEndpoint(token), API_TIMEOUT, EVENTS_API_TIMEOUT, token); RealtimeAPI realtimeApi = new RealtimeAPI(token); - // make api calls on start - Optional res = api.report(Started.get()); - res.ifPresent(ServiceConfigStore::updateFromAPIResponse); - // Fetch blocked lists using separate API call - Optional listsRes = api.fetchBlockedLists(); - listsRes.ifPresent(ServiceConfigStore::updateFromAPIListsResponse); - - // Schedule tasks using ScheduledExecutorService scheduler.scheduleAtFixedRate(new RealtimeTask(realtimeApi, api), POLLING_INTERVAL_SECONDS, POLLING_INTERVAL_SECONDS, TimeUnit.SECONDS); scheduler.scheduleAtFixedRate(new AttackQueueConsumerTask(api), 0, 2, TimeUnit.SECONDS); @@ -60,6 +52,13 @@ public void run() { // Start heartbeats scheduler.schedule(createHeartbeatTask(api, scheduler, true), FIRST_HEARTBEAT_INTERVAL_SECONDS, TimeUnit.SECONDS); + // make api calls on start, slow responses should not delay the tasks above + Optional res = api.report(Started.get()); + res.ifPresent(ServiceConfigStore::updateFromAPIResponse); + // Fetch blocked lists using separate API call + Optional listsRes = api.fetchBlockedLists(); + listsRes.ifPresent(ServiceConfigStore::updateFromAPIListsResponse); + if (token != null) { if (FeatureFlags.AIKIDO_FEATURE_SSE.isEnabled() || ServiceConfigStore.isRealtimeUpdatesEnabled()) { new RealtimeSSETask(new RealtimeSSEAPI(token), api).start(); diff --git a/agent_api/src/test/java/background/HeartbeatSchedulingTest.java b/agent_api/src/test/java/background/HeartbeatSchedulingTest.java index fbf47f12..57f4cf22 100644 --- a/agent_api/src/test/java/background/HeartbeatSchedulingTest.java +++ b/agent_api/src/test/java/background/HeartbeatSchedulingTest.java @@ -82,6 +82,24 @@ void sendsTwoEarlyHeartbeatsThenUsesConfiguredInterval(long interval, boolean re } } + @Test + void schedulesTasksBeforeSendingStartEvent() { + try (var executors = mockStatic(Executors.class); + var clients = mockConstruction(ReportingApiHTTP.class, (client, context) -> + when(client.report(any())).thenAnswer(invocation -> { + verify(scheduler).scheduleAtFixedRate(any(RealtimeTask.class), eq(60L), eq(60L), eq(TimeUnit.SECONDS)); + verify(scheduler).scheduleAtFixedRate(any(AttackQueueConsumerTask.class), eq(0L), eq(2L), eq(TimeUnit.SECONDS)); + verify(scheduler).schedule(any(Runnable.class), eq(30L), eq(TimeUnit.SECONDS)); + return Optional.empty(); + }))) { + executors.when(() -> Executors.newScheduledThreadPool(3)).thenReturn(scheduler); + + new BackgroundProcess("startup-order-test", new Token("token")).run(); + + verify(clients.constructed().get(0)).report(any()); + } + } + @Test void keepsSecondHeartbeatDelayWhenFirstResponseUpdatesInterval() { when(api.report(any())).thenReturn(Optional.of(config("\"heartbeatIntervalInMS\":300000")));