From 02b3d248652172968b2a47432951d8f4a1d337bd Mon Sep 17 00:00:00 2001 From: aktoboy Date: Wed, 1 Nov 2023 15:48:31 +0530 Subject: [PATCH] Added cron to optimize storage --- .../src/main/java/com/akto/testing/Main.java | 2 + .../com/akto/testing/OptimizeStorageCron.java | 99 +++++++++++++++++++ .../testing/TestingOriginalMessageDao.java | 19 ++++ .../dto/testing/TestingOriginalMessage.java | 39 ++++++++ 4 files changed, 159 insertions(+) create mode 100644 apps/testing/src/main/java/com/akto/testing/OptimizeStorageCron.java create mode 100644 libs/dao/src/main/java/com/akto/dao/testing/TestingOriginalMessageDao.java create mode 100644 libs/dao/src/main/java/com/akto/dto/testing/TestingOriginalMessage.java diff --git a/apps/testing/src/main/java/com/akto/testing/Main.java b/apps/testing/src/main/java/com/akto/testing/Main.java index 45386d20c7..d252cd03b5 100644 --- a/apps/testing/src/main/java/com/akto/testing/Main.java +++ b/apps/testing/src/main/java/com/akto/testing/Main.java @@ -112,6 +112,8 @@ public static void main(String[] args) throws InterruptedException { } while (!connectedToMongo); setupRateLimitWatcher(); + OptimizeStorageCron osc = new OptimizeStorageCron(); + osc.init(); loggerMaker.infoAndAddToDb("Starting.......", LogDb.TESTING); diff --git a/apps/testing/src/main/java/com/akto/testing/OptimizeStorageCron.java b/apps/testing/src/main/java/com/akto/testing/OptimizeStorageCron.java new file mode 100644 index 0000000000..78c35dbb1b --- /dev/null +++ b/apps/testing/src/main/java/com/akto/testing/OptimizeStorageCron.java @@ -0,0 +1,99 @@ +package com.akto.testing; + +import com.akto.dao.AccountsDao; +import com.akto.dao.context.Context; +import com.akto.dao.testing.TestingOriginalMessageDao; +import com.akto.dao.testing.TestingRunResultDao; +import com.akto.dto.Account; +import com.akto.dto.ApiInfo; +import com.akto.dto.testing.TestResult; +import com.akto.dto.testing.TestingOriginalMessage; +import com.akto.dto.testing.TestingRunResult; +import com.akto.log.LoggerMaker; +import com.mongodb.BasicDBObject; +import com.mongodb.client.model.Filters; +import org.apache.commons.lang3.tuple.Triple; +import org.bson.types.ObjectId; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static com.akto.dao.MCollection.clients; + +public class OptimizeStorageCron { + + private static final LoggerMaker logger = new LoggerMaker(OptimizeStorageCron.class); + + public void init(){ + List accountList = AccountsDao.instance.findAll(new BasicDBObject(), new BasicDBObject("_id", 1)); + for (Account account : accountList) { + try { + logger.infoAndAddToDb("Starting optimize storage for account: " + account.getId(), LoggerMaker.LogDb.TESTING); + int skip = 0; + int limit = 1000; + Context.accountId.set(account.getId()); + List testingRunResults = TestingRunResultDao.instance.findAll(new BasicDBObject(), skip, limit, new BasicDBObject("_id", -1)); + while (!testingRunResults.isEmpty()) { + logger.infoAndAddToDb("Processing testing run results from: " + skip + " to: " + (skip + limit), LoggerMaker.LogDb.TESTING); + Map> urlToOriginalMessageMap = new HashMap<>(); + for (TestingRunResult testingRunResult : testingRunResults) { + ApiInfo.ApiInfoKey apiInfoKey = testingRunResult.getApiInfoKey(); + ObjectId testingRunResultSummaryId = testingRunResult.getTestRunResultSummaryId(); + List testResults = testingRunResult.getTestResults(); + if (testResults.isEmpty()) { + logger.infoAndAddToDb("No test result found for testing run result: " + testingRunResult.getId(), LoggerMaker.LogDb.TESTING); + continue; + } + int fixedCount = 0; + int notFixedCount = 0; + for (TestResult testResult : testResults) { + if (testResult.getOriginalMessage() != null) { + notFixedCount++; + testResult.setOriginalMessage(null); + urlToOriginalMessageMap.putIfAbsent(apiInfoKey.getUrl(), Triple.of(apiInfoKey, testingRunResultSummaryId, testResult.getOriginalMessage())); + } else { + fixedCount++; + } + } + logger.infoAndAddToDb("Fixed count: " + fixedCount + " Not fixed count: " + notFixedCount, LoggerMaker.LogDb.TESTING); + testingRunResult.setTestResults(testResults); + if (notFixedCount > 0) { + TestingRunResultDao.instance.replaceOne(Filters.eq("_id", testingRunResult.getId()), testingRunResult); + } + } + if (!urlToOriginalMessageMap.isEmpty()) { + for (Map.Entry> entry : urlToOriginalMessageMap.entrySet()) { + Triple triple = entry.getValue(); + ApiInfo.ApiInfoKey apiInfoKey = triple.getLeft(); + ObjectId testingRunResultSummaryId = triple.getMiddle(); + String originalMessage = triple.getRight(); + TestingOriginalMessage testingOriginalMessage = TestingOriginalMessageDao.instance.findOne(Filters.and(Filters.eq(TestingOriginalMessage.API_INFO_KEY, apiInfoKey), Filters.eq(TestingOriginalMessage.TESTING_RUN_RESULT_SUMMARY_ID, testingRunResultSummaryId))); + if (testingOriginalMessage != null) { + logger.infoAndAddToDb("Original message already exists for url: " + apiInfoKey.getUrl() + " method: " + apiInfoKey.getMethod() + " testingRunResultSummaryId: " + testingRunResultSummaryId, LoggerMaker.LogDb.TESTING); + continue; + } + logger.infoAndAddToDb("Inserting original message for url: " + apiInfoKey.getUrl() + " method: " + apiInfoKey.getMethod() + " testingRunResultSummaryId: " + testingRunResultSummaryId, LoggerMaker.LogDb.TESTING); + testingOriginalMessage = new TestingOriginalMessage(); + testingOriginalMessage.setOriginalMessage(originalMessage); + testingOriginalMessage.setApiInfoKey(apiInfoKey); + testingOriginalMessage.setTestingRunResultSummaryId(testingRunResultSummaryId); + TestingOriginalMessageDao.instance.insertOne(testingOriginalMessage); + } + } + skip = skip + limit; + testingRunResults = TestingRunResultDao.instance.findAll(new BasicDBObject(), skip, limit, new BasicDBObject("_id", 1)); + } + int accountId = account.getId(); + logger.infoAndAddToDb("Starting compact for account: " + accountId, LoggerMaker.LogDb.TESTING); + int now = Context.now(); + clients[0].getDatabase(String.valueOf(accountId)).runCommand(new BasicDBObject("compact", "testing_run_result")); + int compactTime = Context.now() - now; + logger.infoAndAddToDb("Finished optimizing storage, compact time for account: " + accountId + " is: " + compactTime, LoggerMaker.LogDb.TESTING); + } catch(Exception e) { + e.printStackTrace(); + logger.errorAndAddToDb("Error while optimizing storage for account: " + account.getId(), LoggerMaker.LogDb.TESTING); + } + } + } +} diff --git a/libs/dao/src/main/java/com/akto/dao/testing/TestingOriginalMessageDao.java b/libs/dao/src/main/java/com/akto/dao/testing/TestingOriginalMessageDao.java new file mode 100644 index 0000000000..4587be9169 --- /dev/null +++ b/libs/dao/src/main/java/com/akto/dao/testing/TestingOriginalMessageDao.java @@ -0,0 +1,19 @@ +package com.akto.dao.testing; + +import com.akto.dao.AccountsContextDao; +import com.akto.dto.testing.TestingOriginalMessage; + +public class TestingOriginalMessageDao extends AccountsContextDao { + + public static final TestingOriginalMessageDao instance = new TestingOriginalMessageDao(); + + @Override + public String getCollName() { + return "testing_original_message"; + } + + @Override + public Class getClassT() { + return TestingOriginalMessage.class; + } +} diff --git a/libs/dao/src/main/java/com/akto/dto/testing/TestingOriginalMessage.java b/libs/dao/src/main/java/com/akto/dto/testing/TestingOriginalMessage.java new file mode 100644 index 0000000000..ed40045076 --- /dev/null +++ b/libs/dao/src/main/java/com/akto/dto/testing/TestingOriginalMessage.java @@ -0,0 +1,39 @@ +package com.akto.dto.testing; + +import com.akto.dto.ApiInfo; +import org.bson.types.ObjectId; + +public class TestingOriginalMessage { + private String originalMessage; + + private ObjectId testingRunResultSummaryId; + + public static final String TESTING_RUN_RESULT_SUMMARY_ID = "testingRunResultSummaryId"; + private ApiInfo.ApiInfoKey apiInfoKey; + + public static final String API_INFO_KEY = "apiInfoKey"; + + public String getOriginalMessage() { + return originalMessage; + } + + public void setOriginalMessage(String originalMessage) { + this.originalMessage = originalMessage; + } + + public ObjectId getTestingRunResultSummaryId() { + return testingRunResultSummaryId; + } + + public void setTestingRunResultSummaryId(ObjectId testingRunResultSummaryId) { + this.testingRunResultSummaryId = testingRunResultSummaryId; + } + + public ApiInfo.ApiInfoKey getApiInfoKey() { + return apiInfoKey; + } + + public void setApiInfoKey(ApiInfo.ApiInfoKey apiInfoKey) { + this.apiInfoKey = apiInfoKey; + } +}