fix: issues with job scheduler

Author: abhijeet-akto

.gitignore

@@ -26,4 +26,5 @@ https:
 model.onnx
 test_scripts/*
 run-master.sh
-run-threat-consumer.sh
+run-threat-consumer.sh
+run-account-jobs.sh

apps/account-job-executor/src/main/java/com/akto/account_job_executor/client/CyborgApiClient.java

@@ -1,20 +1,22 @@
 package com.akto.account_job_executor.client;
 
-import com.akto.dto.OriginalHttpRequest;
-import com.akto.dto.OriginalHttpResponse;
 import com.akto.dto.jobs.AccountJob;
 import com.akto.dto.jobs.JobStatus;
 import com.akto.dto.jobs.ScheduleType;
 import com.akto.log.LoggerMaker;
-import com.akto.testing.ApiExecutor;
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.mongodb.BasicDBObject;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.entity.StringEntity;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.apache.http.impl.client.HttpClients;
+import org.apache.http.util.EntityUtils;
 import org.bson.types.ObjectId;
 
-import java.util.Collections;
+import java.io.IOException;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
 /**
@@ -39,21 +41,51 @@ private static String buildCyborgUrl() {
         if (dbAbsHost.endsWith("/")) {
             dbAbsHost = dbAbsHost.substring(0, dbAbsHost.length() - 1);
         }
-        logger.info("Cyborg URL configured: {}", dbAbsHost + "/api");
+        // Don't use logger during static initialization - can cause NullPointerException
+        System.out.println("[CyborgApiClient] Cyborg URL configured: " + dbAbsHost + "/api");
         return dbAbsHost + "/api";
     }
 
     /**
-     * Build HTTP headers with JWT authentication token.
+     * Get JWT authentication token from environment.
      */
-    private static Map<String, List<String>> buildHeaders() {
+    private static String getAuthToken() {
         String token = System.getenv("DATABASE_ABSTRACTOR_SERVICE_TOKEN");
         if (token == null || token.isEmpty()) {
             throw new IllegalStateException("DATABASE_ABSTRACTOR_SERVICE_TOKEN environment variable not set");
         }
-        Map<String, List<String>> headers = new HashMap<>();
-        headers.put("Authorization", Collections.singletonList(token));
-        return headers;
+        return token;
+    }
+
+    /**
+     * Make HTTP POST request to Cyborg API using Apache HttpClient.
+     * This bypasses ApiExecutor's SSRF protection since we're making trusted internal calls.
+     */
+    private static String makePostRequest(String endpoint, Map<String, Object> requestBody) throws IOException {
+        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
+            HttpPost httpPost = new HttpPost(url + endpoint);
+
+            // Set headers
+            httpPost.setHeader("Authorization", getAuthToken());
+            httpPost.setHeader("Content-Type", "application/json");
+
+            // Set body
+            String jsonBody = mapper.writeValueAsString(requestBody);
+            httpPost.setEntity(new StringEntity(jsonBody));
+
+            // Execute request
+            try (CloseableHttpResponse response = httpClient.execute(httpPost)) {
+                int statusCode = response.getStatusLine().getStatusCode();
+                String responseBody = EntityUtils.toString(response.getEntity());
+
+                if (statusCode != 200) {
+                    logger.error("API request failed. Status: {}, Body: {}", statusCode, responseBody);
+                    throw new IOException("API request failed with status: " + statusCode);
+                }
+
+                return responseBody;
+            }
+        }
     }
 
     /**
@@ -65,29 +97,16 @@ private static Map<String, List<String>> buildHeaders() {
      * @return Claimed AccountJob or null if no jobs available
      */
     public static AccountJob fetchAndClaimJob(long now, int heartbeatThreshold) {
+        logger.info("=== fetchAndClaimJob called: now={}, threshold={} ===", now, heartbeatThreshold);
         try {
-            Map<String, List<String>> headers = buildHeaders();
-            BasicDBObject requestBody = new BasicDBObject();
+            Map<String, Object> requestBody = new HashMap<>();
             requestBody.put("now", now);
             requestBody.put("heartbeatThreshold", heartbeatThreshold);
 
-            OriginalHttpRequest request = new OriginalHttpRequest(
-                url + "/fetchAndClaimAccountJob",
-                "", "POST",
-                requestBody.toString(),
-                headers, ""
-            );
-
-            logger.debug("Fetching and claiming job from Cyborg API");
-            OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
+            logger.info("Fetching and claiming job from Cyborg API: now={}, threshold={}", now, heartbeatThreshold);
+            String body = makePostRequest("/fetchAndClaimAccountJob", requestBody);
 
-            if (response.getStatusCode() != 200) {
-                logger.error("Failed to fetch job. Status: {}, Body: {}",
-                    response.getStatusCode(), response.getBody());
-                return null;
-            }
-
-            String body = response.getBody();
+            logger.info("API response received, body length: {}", body != null ? body.length() : 0);
             if (body == null || body.isEmpty() || body.equals("null") || body.equals("{}")) {
                 logger.debug("No jobs available to claim");
                 return null;
@@ -107,7 +126,8 @@ public static AccountJob fetchAndClaimJob(long now, int heartbeatThreshold) {
             return job;
 
         } catch (Exception e) {
-            logger.error("Error fetching and claiming job", e);
+            logger.error("!!! EXCEPTION in fetchAndClaimJob: {}", e.getMessage(), e);
+            e.printStackTrace();
             return null;
         }
     }
@@ -123,31 +143,17 @@ public static AccountJob fetchAndClaimJob(long now, int heartbeatThreshold) {
      */
     public static void updateJobStatus(ObjectId id, JobStatus status, Integer finishedAt, String error) {
         try {
-            Map<String, List<String>> headers = buildHeaders();
-            BasicDBObject requestBody = new BasicDBObject();
+            Map<String, Object> requestBody = new HashMap<>();
             requestBody.put("jobId", id.toString());
             requestBody.put("status", status.name());
             requestBody.put("finishedAt", finishedAt);
             if (error != null) {
                 requestBody.put("error", error);
             }
 
-            OriginalHttpRequest request = new OriginalHttpRequest(
-                url + "/updateAccountJobStatus",
-                "", "POST",
-                requestBody.toString(),
-                headers, ""
-            );
-
             logger.debug("Updating job status: jobId={}, status={}", id, status);
-            OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
-
-            if (response.getStatusCode() != 200) {
-                logger.error("Failed to update job status. Status: {}, Body: {}",
-                    response.getStatusCode(), response.getBody());
-            } else {
-                logger.debug("Successfully updated job status: jobId={}, status={}", id, status);
-            }
+            makePostRequest("/updateAccountJobStatus", requestBody);
+            logger.debug("Successfully updated job status: jobId={}, status={}", id, status);
 
         } catch (Exception e) {
             logger.error("Error updating job status", e);
@@ -162,26 +168,12 @@ public static void updateJobStatus(ObjectId id, JobStatus status, Integer finish
      */
     public static void updateJobHeartbeat(ObjectId id) {
         try {
-            Map<String, List<String>> headers = buildHeaders();
-            BasicDBObject requestBody = new BasicDBObject();
+            Map<String, Object> requestBody = new HashMap<>();
             requestBody.put("jobId", id.toString());
 
-            OriginalHttpRequest request = new OriginalHttpRequest(
-                url + "/updateAccountJobHeartbeat",
-                "", "POST",
-                requestBody.toString(),
-                headers, ""
-            );
-
             logger.debug("Updating job heartbeat: jobId={}", id);
-            OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
-
-            if (response.getStatusCode() != 200) {
-                logger.error("Failed to update job heartbeat. Status: {}, Body: {}",
-                    response.getStatusCode(), response.getBody());
-            } else {
-                logger.debug("Successfully updated job heartbeat: jobId={}", id);
-            }
+            makePostRequest("/updateAccountJobHeartbeat", requestBody);
+            logger.debug("Successfully updated job heartbeat: jobId={}", id);
 
         } catch (Exception e) {
             logger.error("Error updating job heartbeat", e);
@@ -197,27 +189,13 @@ public static void updateJobHeartbeat(ObjectId id) {
      */
     public static void updateJob(ObjectId id, Map<String, Object> updates) {
         try {
-            Map<String, List<String>> headers = buildHeaders();
-            BasicDBObject requestBody = new BasicDBObject();
+            Map<String, Object> requestBody = new HashMap<>();
             requestBody.put("jobId", id.toString());
             requestBody.put("updates", updates);
 
-            OriginalHttpRequest request = new OriginalHttpRequest(
-                url + "/updateAccountJob",
-                "", "POST",
-                requestBody.toString(),
-                headers, ""
-            );
-
             logger.debug("Updating job: jobId={}, updates={}", id, updates);
-            OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
-
-            if (response.getStatusCode() != 200) {
-                logger.error("Failed to update job. Status: {}, Body: {}",
-                    response.getStatusCode(), response.getBody());
-            } else {
-                logger.debug("Successfully updated job: jobId={}", id);
-            }
+            makePostRequest("/updateAccountJob", requestBody);
+            logger.debug("Successfully updated job: jobId={}", id);
 
         } catch (Exception e) {
             logger.error("Error updating job", e);
@@ -232,27 +210,11 @@ public static void updateJob(ObjectId id, Map<String, Object> updates) {
      */
     public static AccountJob findById(ObjectId id) {
         try {
-            Map<String, List<String>> headers = buildHeaders();
-            BasicDBObject requestBody = new BasicDBObject();
+            Map<String, Object> requestBody = new HashMap<>();
             requestBody.put("jobId", id.toString());
 
-            OriginalHttpRequest request = new OriginalHttpRequest(
-                url + "/fetchAccountJob",
-                "", "POST",
-                requestBody.toString(),
-                headers, ""
-            );
-
             logger.debug("Fetching job by ID: jobId={}", id);
-            OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
-
-            if (response.getStatusCode() != 200) {
-                logger.error("Failed to fetch job. Status: {}, Body: {}",
-                    response.getStatusCode(), response.getBody());
-                return null;
-            }
-
-            String body = response.getBody();
+            String body = makePostRequest("/fetchAccountJob", requestBody);
             if (body == null || body.isEmpty() || body.equals("null") || body.equals("{}")) {
                 logger.debug("Job not found: jobId={}", id);
                 return null;
@@ -286,14 +248,25 @@ public static AccountJob findById(ObjectId id) {
     private static AccountJob convertMapToAccountJob(Map<String, Object> map) {
         AccountJob job = new AccountJob();
 
-        // Handle ObjectId
-        if (map.get("_id") != null) {
-            Object idObj = map.get("_id");
+        // Handle ObjectId - check both "_id" and "id" fields
+        Object idObj = map.get("_id");
+        if (idObj == null) {
+            idObj = map.get("id");  // Try "id" without underscore
+        }
+
+        if (idObj != null) {
             if (idObj instanceof String) {
                 job.setId(new ObjectId((String) idObj));
             } else if (idObj instanceof Map) {
                 Map<String, Object> idMap = (Map<String, Object>) idObj;
-                job.setId(new ObjectId((String) idMap.get("$oid")));
+                // Handle different ObjectId formats: {"$oid": "..."} or {"timestamp": ..., "date": ...}
+                if (idMap.containsKey("$oid")) {
+                    job.setId(new ObjectId((String) idMap.get("$oid")));
+                } else if (idMap.containsKey("timestamp")) {
+                    // Create ObjectId from timestamp
+                    int timestamp = getIntValue(idMap, "timestamp");
+                    job.setId(new ObjectId(timestamp, 0));
+                }
             }
         }
 

apps/account-job-executor/src/main/java/com/akto/account_job_executor/cron/AccountJobsCron.java

@@ -21,7 +21,7 @@ public class AccountJobsCron {
 
     private static final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
     private static final ExecutorService executorService = Executors.newFixedThreadPool(2);
-    private static final int MAX_HEARTBEAT_THRESHOLD_SECONDS = 300;
+    private static final int MAX_HEARTBEAT_THRESHOLD_SECONDS = 5;
 
     private static final LoggerMaker logger = new LoggerMaker(AccountJobsCron.class);
 
@@ -39,16 +39,19 @@ public void startScheduler() {
         scheduler.scheduleAtFixedRate(() -> {
             try {
                 long now = Context.now();
+                logger.info("Polling for jobs: now={}, threshold={}", now, MAX_HEARTBEAT_THRESHOLD_SECONDS);
 
                 // Fetch and claim one job atomically via Cyborg API
                 // This returns either:
                 // 1. A SCHEDULED job where scheduledAt < now
                 // 2. A RUNNING job where heartbeatAt < now - 300s (stale job)
                 // 3. null if no jobs available
+                logger.info("About to call CyborgApiClient.fetchAndClaimJob()...");
                 AccountJob job = CyborgApiClient.fetchAndClaimJob(now, MAX_HEARTBEAT_THRESHOLD_SECONDS);
+                logger.info("CyborgApiClient.fetchAndClaimJob() returned: {}", job != null ? "job found" : "null");
 
                 if (job == null) {
-                    logger.debug("No jobs to run");
+                    logger.info("No jobs to run at this time");
                     return;
                 }
 
@@ -76,7 +79,13 @@ public void startScheduler() {
                 });
 
             } catch (Exception e) {
-                logger.error("Error in AccountJobs scheduler", e);
+                logger.error("!!! EXCEPTION in AccountJobs scheduler !!!", e);
+                e.printStackTrace();
+                System.err.println("EXCEPTION: " + e.getMessage());
+            } catch (Throwable t) {
+                logger.error("!!! THROWABLE in AccountJobs scheduler !!!", t);
+                t.printStackTrace();
+                System.err.println("THROWABLE: " + t.getMessage());
             }
         }, 0, 5, TimeUnit.SECONDS);