Refactoring - adding details of query execution to KMessage

dzmipt · Dec 8, 2023 · bf53f9d · bf53f9d
1 parent 6372874
commit bf53f9d
Show file tree

Hide file tree

Showing 8 changed files with 78 additions and 34 deletions.
diff --git a/src/kx/KConnection.java b/src/kx/KConnection.java
@@ -123,7 +123,7 @@ public void setConnectionStateListener(ConnectionStateListener connectionStateLi
 
     private final static byte[] HEADER = new byte[] {0,1,0,0};
 
-    private void send(K.KBase query) throws IOException {
+    private long send(K.KBase query) throws IOException {
         ByteArrayOutputStream baosBody = new ByteArrayOutputStream();
         query.serialise(baosBody);
 
@@ -134,22 +134,28 @@ private void send(K.KBase query) throws IOException {
 
         outputStream.write(baosHeader.toByteArray());
         outputStream.write(baosBody.toByteArray());
-        stats.sentBytes(baosHeader.size() + baosBody.size());
+        long sentBytes = baosHeader.size() + baosBody.size();
+        stats.sentBytes(sentBytes);
+        return sentBytes;
     }
 
-    public synchronized K.KBase k(K.KBase x, ProgressCallback progress) throws K4Exception, IOException, InterruptedException {
+    public synchronized KMessage k(K.KBase x, ProgressCallback progress) throws K4Exception, IOException, InterruptedException {
         try {
             if (isClosed()) connect();
             socketReader.setProgressCallback(progress);
-            send(x);
-            return socketReader.getResponse();
+            K.KTimestamp sentTime = K.KTimestamp.now();
+            long sentBytes = send(x);
+            KMessage message = socketReader.getResponse();
+            message.setStarted(sentTime);
+            message.setBytesSent(sentBytes);
+            return message;
         } catch (IOException e) {
             close();
             throw e;
         }
     }
 
-    public K.KBase k(K.KBase x) throws K4Exception, IOException, InterruptedException {
+    public KMessage k(K.KBase x) throws K4Exception, IOException, InterruptedException {
         return k(x, null);
     }
 
@@ -174,7 +180,7 @@ private synchronized ProgressCallback getProgressCallback() {
             return progress;
         }
 
-        K.KBase getResponse() throws InterruptedException, K4Exception, IOException {
+        KMessage getResponse() throws InterruptedException, IOException {
             KMessage response;
             synchronized (lockRead) {
                 while (message == null) {
@@ -187,7 +193,7 @@ K.KBase getResponse() throws InterruptedException, K4Exception, IOException {
                 message = null;
                 lockWrite.notifyAll();
             }
-            return response.getObject();
+            return response;
         }
 
         @Override
@@ -208,6 +214,7 @@ public void run() {
                     IPC ipc = new IPC(buffer, 4, false, isLittleEndian);
                     final int msgLength = ipc.ri() - 8;
 
+                    K.KTimestamp receivedTime = K.KTimestamp.now();
                     if (response && connectionStateListener != null) {
                         connectionStateListener.checkIncomingLimit(msgLength);
                     }
@@ -233,6 +240,8 @@ public void run() {
                     if (response) {
                         synchronized (lockRead) {
                             message = IPC.deserialise(buffer, compressed, isLittleEndian);
+                            message.setBytesReceived(msgLength);
+                            message.setFinished(receivedTime);
                             lockRead.notifyAll();
                         }
                         synchronized (lockWrite) {

diff --git a/src/kx/KMessage.java b/src/kx/KMessage.java
@@ -9,6 +9,11 @@ public class KMessage {
     private K.KBase object = null;
     private K4Exception error = null;
     private IOException exception = null;
+    private long bytesSent = 0;
+    private long bytesReceived = 0;
+    private K.KTimestamp started = K.KTimestamp.NULL;
+    private K.KTimestamp finished = K.KTimestamp.NULL;
+
 
     public KMessage(K.KBase result) {
         this.object = result;
@@ -22,10 +27,45 @@ public KMessage(IOException exception) {
         this.exception = exception;
     }
 
-    public K.KBase getObject() throws K4Exception, IOException {
-        if (error != null) throw error;
-        if (exception != null) throw exception;
-
+    public K.KBase getObject() {
         return object;
     }
+
+    public Throwable getError() {
+        if (error != null) return  error;
+        if (exception != null) return  exception;
+        return null;
+    }
+
+    public long getBytesSent() {
+        return bytesSent;
+    }
+
+    public void setBytesSent(long bytesSent) {
+        this.bytesSent = bytesSent;
+    }
+
+    public long getBytesReceived() {
+        return bytesReceived;
+    }
+
+    public void setBytesReceived(long bytesReceived) {
+        this.bytesReceived = bytesReceived;
+    }
+
+    public K.KTimestamp getStarted() {
+        return started;
+    }
+
+    public void setStarted(K.KTimestamp started) {
+        this.started = started;
+    }
+
+    public K.KTimestamp getFinished() {
+        return finished;
+    }
+
+    public void setFinished(K.KTimestamp finished) {
+        this.finished = finished;
+    }
 }
diff --git a/src/studio/kdb/Session.java b/src/studio/kdb/Session.java
@@ -111,7 +111,7 @@ static void mock(SessionCreator sessionCreator) {
         Session.sessionCreator = sessionCreator;
     }
 
-    public K.KBase execute(K.KBase x, ProgressCallback progress) throws K4Exception, IOException, InterruptedException {
+    public KMessage execute(K.KBase x, ProgressCallback progress) throws K4Exception, IOException, InterruptedException {
         return kConn.k(x, progress);
     }
 

diff --git a/src/studio/ui/action/ConnectionStats.java b/src/studio/ui/action/ConnectionStats.java
@@ -1,6 +1,7 @@
 package studio.ui.action;
 
 import kx.KConnectionStats;
+import kx.KMessage;
 import studio.kdb.K;
 import studio.kdb.Session;
 import studio.ui.StudioWindow;
@@ -12,7 +13,7 @@ public class ConnectionStats {
 
     public static void getStats(StudioWindow studioWindow) {
         QueryResult queryResult = new QueryResult(null, "");
-        queryResult.setResult(getTable());
+        queryResult.setResult(new KMessage(getTable()));
         studioWindow.addResultTab(queryResult, "Connection statistics");
     }
 

diff --git a/src/studio/ui/action/QueryExecutor.java b/src/studio/ui/action/QueryExecutor.java
@@ -1,6 +1,7 @@
 package studio.ui.action;
 
 import kx.K4Exception;
+import kx.KMessage;
 import kx.ProgressCallback;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -127,10 +128,9 @@ private Session getSession() {
         protected QueryResult doInBackground() {
             QueryResult result = new QueryResult(server, queryText);
             queryLog.info("#{}: query {}({})\n{}",queryIndex, server.getFullName(), server.getConnectionString(), queryText);
-            long startTime = System.currentTimeMillis();
             try {
                 session = getSession();
-                K.KBase response = session.execute(query, QueryExecutor.this);
+                KMessage response = session.execute(query, QueryExecutor.this);
                 result.setResult(response);
             } catch (Throwable e) {
                 if (! (e instanceof K4Exception)) {
@@ -139,7 +139,6 @@ protected QueryResult doInBackground() {
                 }
                 result.setError(e);
             }
-            result.setExecutionTime(System.currentTimeMillis() - startTime);
             if (result.getError() != null) {
                 if (result.getError() instanceof K4Exception) {
                     queryLog.info("#{}: server returns error {}", queryIndex, result.getError().getMessage());

diff --git a/src/studio/ui/action/QueryResult.java b/src/studio/ui/action/QueryResult.java
@@ -1,5 +1,6 @@
 package studio.ui.action;
 
+import kx.KMessage;
 import studio.kdb.K;
 import studio.kdb.Server;
 
@@ -8,18 +9,19 @@ public class QueryResult {
     private String query;
     private Server server;
 
-    private K.KBase result = null;
+    private KMessage result = null;
     private Throwable error = null;
-    private long executionTime = -1;
     private boolean complete = false;
 
     public QueryResult(Server server, String query) {
         this.server = server;
         this.query = query;
     }
 
-    public void setResult(K.KBase result) {
+    public void setResult(KMessage result) {
         this.result = result;
+        Throwable error = result.getError();
+        if (error != null) setError(error);
         complete = true;
     }
 
@@ -28,10 +30,6 @@ public void setError(Throwable error) {
         complete = true;
     }
 
-    public void setExecutionTime(long executionTime) {
-        this.executionTime = executionTime;
-    }
-
     public boolean isComplete() {
         return complete;
     }
@@ -45,14 +43,14 @@ public Server getServer() {
     }
 
     public K.KBase getResult() {
-        return result;
+        return result.getObject();
     }
 
     public Throwable getError() {
         return error;
     }
 
     public long getExecutionTime() {
-        return executionTime;
+        return result.getFinished().toLong() - result.getStarted().toLong();
     }
 }
diff --git a/test-integration/studio/kdb/MockQSession.java b/test-integration/studio/kdb/MockQSession.java
@@ -1,9 +1,6 @@
 package studio.kdb;
 
-import kx.K4Exception;
-import kx.KAuthentication;
-import kx.KConnection;
-import kx.ProgressCallback;
+import kx.*;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -93,7 +90,7 @@ public boolean isClosed() {
     }
 
     @Override
-    public synchronized K.KBase k(K.KBase x, ProgressCallback progress) throws K4Exception, IOException {
+    public synchronized KMessage k(K.KBase x, ProgressCallback progress) throws K4Exception, IOException {
         log.info("MockQSession.k - query execution");
         queryCount.getAndIncrement();
         closed = false;
@@ -110,7 +107,7 @@ public synchronized K.KBase k(K.KBase x, ProgressCallback progress) throws K4Exc
             log.info("MockQSession.k - continue after locking");
         }
 
-        if (kResponse != null) return kResponse;
+        if (kResponse != null) return new KMessage(kResponse);
         if (kError != null) throw kError;
 
         if (ioException != null) {
@@ -119,7 +116,7 @@ public synchronized K.KBase k(K.KBase x, ProgressCallback progress) throws K4Exc
         }
 
 
-        return x;
+        return new KMessage(x);
     }
 
     public int getQueryCount() {

diff --git a/test/studio/kdb/KSerialiseTest.java b/test/studio/kdb/KSerialiseTest.java
@@ -32,7 +32,7 @@ public static void exit() throws K4Exception, IOException, InterruptedException
 
     private void test(K.KBase k) {
         try {
-            K.KBase result = kConn.k(k);
+            K.KBase result = kConn.k(k).getObject();
             assertEquals(k, result);
         } catch (K4Exception | IOException | InterruptedException e) {
             fail(e);