From 11c192e6c2cda56e912d3155768ffdf84b085b24 Mon Sep 17 00:00:00 2001
From: Gagan Juneja <gjjuneja@amazon.com>
Date: Tue, 7 May 2024 07:58:11 +0530
Subject: [PATCH] test changes

---
 .../listener/PerformanceAnalyzerSearchListener.java   |  2 +-
 .../PerformanceAnalyzerTransportChannel.java          | 11 ++++++++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/opensearch/performanceanalyzer/listener/PerformanceAnalyzerSearchListener.java b/src/main/java/org/opensearch/performanceanalyzer/listener/PerformanceAnalyzerSearchListener.java
index 1c011383..be588c2c 100644
--- a/src/main/java/org/opensearch/performanceanalyzer/listener/PerformanceAnalyzerSearchListener.java
+++ b/src/main/java/org/opensearch/performanceanalyzer/listener/PerformanceAnalyzerSearchListener.java
@@ -149,7 +149,7 @@ protected void innerOnResponse(Task task) {
                                             Tags.create()
                                                     .addTag(
                                                             "shardId",
-                                                            searchContext.getTask().getId())
+                                                            searchContext.shardTarget().getShardId().getId())
                                                     .addTag(
                                                             "indexName",
                                                             searchContext
diff --git a/src/main/java/org/opensearch/performanceanalyzer/transport/PerformanceAnalyzerTransportChannel.java b/src/main/java/org/opensearch/performanceanalyzer/transport/PerformanceAnalyzerTransportChannel.java
index d9dae89c..c944ec7d 100644
--- a/src/main/java/org/opensearch/performanceanalyzer/transport/PerformanceAnalyzerTransportChannel.java
+++ b/src/main/java/org/opensearch/performanceanalyzer/transport/PerformanceAnalyzerTransportChannel.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.performanceanalyzer.transport;
 
+import com.sun.management.ThreadMXBean;
 import java.io.BufferedReader;
 import java.io.File;
 import java.io.FileNotFoundException;
@@ -46,6 +47,7 @@ public class PerformanceAnalyzerTransportChannel implements TransportChannel, Me
     private String threadID;
     private Counter throughputCounter;
     private Long startCpuTimeNanos;
+    private long jvmThreadID;
 
     void set(
             TransportChannel original,
@@ -64,6 +66,8 @@ void set(
         this.startValue = getIOUsage(threadID);
         this.throughputCounter = throughputCounter;
         this.startCpuTimeNanos = ManagementFactory.getThreadMXBean().getCurrentThreadCpuTime();
+        jvmThreadID = Thread.currentThread().getId();
+        LOG.info("Updating the counter Start thread {} {}", jvmThreadID, Thread.currentThread().getName());
     }
 
     @Override
@@ -78,10 +82,11 @@ public String getChannelType() {
 
     @Override
     public void sendResponse(TransportResponse response) throws IOException {
+        LOG.info("Updating the counter Finish thread {}", Thread.currentThread().getName());
         if (throughputCounter != null) {
-            long updatedValue = ManagementFactory.getThreadMXBean().getCurrentThreadCpuTime()
-                    - startCpuTimeNanos;
-            LOG.info("Updating the counter inside PATransportChannel {}", updatedValue);
+            long timeNow = ((ThreadMXBean)ManagementFactory.getThreadMXBean()).getThreadCpuTime(jvmThreadID);
+            long updatedValue = timeNow - startCpuTimeNanos;
+            LOG.info("Updating the counter inside PATransportChannel {} start time {}, now {}", updatedValue, startCpuTimeNanos, timeNow);
             throughputCounter.add(
                     Math.max(updatedValue, 0),
                     Tags.create().addTag("indexName", indexName).addTag("shardId", shardId));