APM instrumentation improvements

jadevogt · Sep 17, 2023 · 91f3b14 · 91f3b14
1 parent b8a5685
commit 91f3b14
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 38 deletions.
diff --git a/pom.xml b/pom.xml
@@ -6,7 +6,7 @@
 
     <groupId>best.tigers</groupId>
     <artifactId>tigersbot</artifactId>
-    <version>1.3.0</version>
+    <version>1.3.1</version>
     <properties>
         <maven.compiler.source>19</maven.compiler.source>
         <maven.compiler.target>19</maven.compiler.target>

diff --git a/src/main/java/best/tigers/tigersbot/services/BotService.java b/src/main/java/best/tigers/tigersbot/services/BotService.java
@@ -56,7 +56,7 @@ private void processUpdate(Update update) {
             if (handler.invokationTest(message)) {
                 Transaction transaction = ElasticApm.startTransaction();
                 try {
-                    transaction.setName(handler.getClass().getName());
+                    transaction.setName(handler.getClass().getSimpleName());
                     transaction.setType(Transaction.TYPE_REQUEST);
                     handler.handle(message);
                 } catch (Throwable e) {

diff --git a/src/main/java/best/tigers/tigersbot/services/CompletionService.java b/src/main/java/best/tigers/tigersbot/services/CompletionService.java
@@ -3,6 +3,7 @@
 import best.tigers.tigersbot.error.MissingEnvironmentVariableException;
 import best.tigers.tigersbot.util.Log;
 import co.elastic.apm.api.CaptureSpan;
+import co.elastic.apm.api.ElasticApm;
 import com.theokanning.openai.completion.CompletionRequest;
 import com.theokanning.openai.completion.chat.ChatCompletionRequest;
 import com.theokanning.openai.completion.chat.ChatMessage;
@@ -92,18 +93,25 @@ public String getCompletion(String prompt, String modelName, String userIdentifi
             " at least once every other sentence."
     );
 
-    @CaptureSpan(type = "external", subtype = "openai", action = "chatcompletion")
     public String getAdvancedCompletion(String prompt, String userIdentifier, ChatMessage systemMessage, String model) {
-        var userChatMessage = new ChatMessage(
-                ChatMessageRole.USER.value(),
-                prompt
-        );
-        var advancedCompletionRequest = ChatCompletionRequest.builder()
-                .messages(List.of(systemMessage, userChatMessage))
-                .model(model)
-                .user(String.valueOf(userIdentifier.hashCode()))
-                .build();
-        return api.createChatCompletion(advancedCompletionRequest).getChoices().get(0).getMessage().getContent();
+        var span = ElasticApm.currentTransaction().startSpan("external", "openai", "chatcompletion");
+        try {
+            var userChatMessage = new ChatMessage(
+                    ChatMessageRole.USER.value(),
+                    prompt
+            );
+            var advancedCompletionRequest = ChatCompletionRequest.builder()
+                    .messages(List.of(systemMessage, userChatMessage))
+                    .model(model)
+                    .user(String.valueOf(userIdentifier.hashCode()))
+                    .build();
+            return api.createChatCompletion(advancedCompletionRequest).getChoices().get(0).getMessage().getContent();
+        } catch (Throwable e) {
+            span.captureException(e);
+            throw e;
+        } finally {
+            span.end();
+        }
     }
 
     public String getAdvancedCompletion(String prompt, String userIdentifier, ChatMessage systemMessage) {
@@ -115,14 +123,21 @@ public String getAdvancedCompletion(String prompt, String userIdentifier) {
     }
 
 
-    @CaptureSpan(type = "external", subtype = "openai", action = "imagecompletion")
     public Image getImage(String prompt) {
-        var completionRequest = CreateImageRequest.builder()
-                .prompt(prompt)
-                .n(1)
-                .size("1024x1024")
-                .build();
-        return api.createImage(completionRequest).getData().get(0);
+        var span = ElasticApm.currentTransaction().startSpan("external", "openai", "imagecompletion");
+        try {
+            var completionRequest = CreateImageRequest.builder()
+                    .prompt(prompt)
+                    .n(1)
+                    .size("1024x1024")
+                    .build();
+            return api.createImage(completionRequest).getData().get(0);
+        } catch (Throwable e) {
+            span.captureException(e);
+            throw e;
+        } finally {
+            span.end();
+        }
     }
 
 

diff --git a/src/main/java/best/tigers/tigersbot/services/ImagesService.java b/src/main/java/best/tigers/tigersbot/services/ImagesService.java
@@ -4,6 +4,7 @@
 import best.tigers.tigersbot.util.Environment;
 import best.tigers.tigersbot.util.Log;
 import co.elastic.apm.api.CaptureSpan;
+import co.elastic.apm.api.ElasticApm;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -32,13 +33,20 @@ public static ImagesService getInstance() throws MissingEnvironmentVariableExcep
         return instance;
     }
 
-    @CaptureSpan(type = "external", subtype = "googleimages", action = "query")
     public String getImage(String query) {
-        var conn = getConnection(query);
-        if (conn == null) {
-            return "";
+        var span = ElasticApm.currentTransaction().startSpan("external", "googleimages", "query");
+        try {
+            var conn = getConnection(query);
+            if (conn == null) {
+                return "";
+            }
+            return getImageLink(conn);
+        } catch (Throwable e) {
+            span.captureException(e);
+            throw e;
+        } finally {
+            span.end();
         }
-        return getImageLink(conn);
     }
 
     private HttpURLConnection getConnection(String query) {

diff --git a/src/main/java/best/tigers/tigersbot/services/PersistentStorageService.java b/src/main/java/best/tigers/tigersbot/services/PersistentStorageService.java
@@ -2,6 +2,7 @@
 
 import best.tigers.tigersbot.util.Log;
 import co.elastic.apm.api.CaptureSpan;
+import co.elastic.apm.api.ElasticApm;
 import org.json.JSONObject;
 import org.json.JSONTokener;
 
@@ -19,24 +20,39 @@ public static PersistentStorageService getInstance() {
         return instance;
     }
 
-    @CaptureSpan(type = "io", subtype = "filesystem", action = "read")
     public synchronized JSONObject jsonFromFile(String filePath) {
-        JSONObject jsonObject = null;
-        try (FileReader reader = new FileReader(filePath)) {
-            jsonObject = new JSONObject(new JSONTokener(reader));
-        } catch (IOException e) {
-            Log.severe(e.getMessage());
+        var span = ElasticApm.currentTransaction().startSpan("io", "filesystem", "read");
+        try {
+            JSONObject jsonObject = null;
+            try (FileReader reader = new FileReader(filePath)) {
+                jsonObject = new JSONObject(new JSONTokener(reader));
+            } catch (IOException e) {
+                Log.severe(e.getMessage());
+            }
+            return jsonObject;
+        } catch (Throwable e) {
+            span.captureException(e);
+            throw e;
+        } finally {
+            span.end();
         }
-        return jsonObject;
     }
 
-    @CaptureSpan(type = "io", subtype = "filesystem", action = "write")
     public synchronized void jsonToFile(String filePath, JSONObject jsonObject) {
-        try (FileWriter writer = new FileWriter(filePath)) {
-            writer.write(jsonObject.toString());
-            writer.flush();
-        } catch (Exception e) {
-            Log.severe(e.getMessage());
+        var span = ElasticApm.currentTransaction().startSpan("io", "filesystem", "write");
+        try {
+            try (FileWriter writer = new FileWriter(filePath)) {
+                writer.write(jsonObject.toString());
+                writer.flush();
+            } catch (Exception e) {
+                Log.severe(e.getMessage());
+            }
+        } catch (Throwable e) {
+            span.captureException(e);
+            throw e;
+        } finally {
+            span.end();
         }
     }
+
 }