Set DocumentOperationType in ReplyContext

4 files changed, 20 insertions, 15 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java b/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java
index 56849b9db8f..8f699c156b4 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java
@@ -1,25 +1,26 @@
 // Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.documentapi.metrics;
 
-import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
+import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.RemoveDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.UpdateDocumentMessage;
 import com.yahoo.messagebus.Message;
 
 /**
  * @author freva
  */
 public enum DocumentOperationType {
-    PUT, REMOVE, UPDATE;
+    PUT, REMOVE, UPDATE, ERROR;
 
     public static DocumentOperationType fromMessage(Message msg) {
-        switch (msg.getType()) {
-            case DocumentProtocol.MESSAGE_PUTDOCUMENT:
-                return PUT;
-            case DocumentProtocol.MESSAGE_UPDATEDOCUMENT:
-                return UPDATE;
-            case DocumentProtocol.MESSAGE_REMOVEDOCUMENT:
-                return REMOVE;
-            default:
-                return null;
+        if (msg instanceof PutDocumentMessage) {
+            return PUT;
+        } else if (msg instanceof RemoveDocumentMessage) {
+            return REMOVE;
+        } else if (msg instanceof UpdateDocumentMessage) {
+            return UPDATE;
+        } else {
+            return ERROR;
         }
     }
 }
diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ClientFeederV3.java b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ClientFeederV3.java
index f91038a0a79..8c1344efca1 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ClientFeederV3.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ClientFeederV3.java
@@ -6,6 +6,7 @@ import com.yahoo.container.jdisc.HttpResponse;
 import com.yahoo.document.DocumentTypeManager;
 import com.yahoo.documentapi.messagebus.protocol.DocumentMessage;
 import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
+import com.yahoo.documentapi.metrics.DocumentOperationType;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.jdisc.ReferencedResource;
 import com.yahoo.log.LogLevel;
@@ -301,7 +302,7 @@ class ClientFeederV3 {
     }
 
     private void setMessageParameters(DocumentOperationMessageV3 msg, FeederSettings settings) {
-        msg.getMessage().setContext(new ReplyContext(msg.getOperationId(), feedReplies));
+        msg.getMessage().setContext(new ReplyContext(msg.getOperationId(), feedReplies, DocumentOperationType.fromMessage(msg.getMessage())));
         if (settings.traceLevel != null) {
             msg.getMessage().getTrace().setLevel(settings.traceLevel);
         }
diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/Feeder.java b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/Feeder.java
index ad773519458..ff7afb26410 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/Feeder.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/Feeder.java
@@ -13,6 +13,7 @@ import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
 import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
 import com.yahoo.documentapi.messagebus.protocol.RemoveDocumentMessage;
 import com.yahoo.documentapi.messagebus.protocol.UpdateDocumentMessage;
+import com.yahoo.documentapi.metrics.DocumentOperationType;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.jdisc.ReferencedResource;
 import com.yahoo.log.LogLevel;
@@ -350,7 +351,7 @@ public class Feeder implements Runnable {
     }
 
     private void setMessageParameters(Tuple2<String, Message> msg) {
-        msg.second.setContext(new ReplyContext(msg.first, feedReplies));
+        msg.second.setContext(new ReplyContext(msg.first, feedReplies, DocumentOperationType.fromMessage(msg.second)));
         if (settings.traceLevel != null) {
             msg.second.getTrace().setLevel(settings.traceLevel);
         }
diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ReplyContext.java b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ReplyContext.java
index 30d8b04131d..39392fdefc4 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ReplyContext.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/ReplyContext.java
@@ -1,6 +1,7 @@
 // Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.http.server;
 
+import com.yahoo.documentapi.metrics.DocumentOperationType;
 import com.yahoo.vespa.http.client.core.OperationStatus;
 
 import java.util.concurrent.BlockingQueue;
@@ -8,15 +9,16 @@ import java.util.concurrent.BlockingQueue;
 /**
  * Mapping between document ID and client session.
  *
- * @author <a href="mailto:steinar@yahoo-inc.com">Steinar Knutsen</a>
+ * @author Steinar Knutsen
  */
 public class ReplyContext {
 
     public final String docId;
+    public DocumentOperationType documentOperationType;
     public final BlockingQueue<OperationStatus> feedReplies;
     public final long creationTime;
 
-    public ReplyContext(String docId, BlockingQueue<OperationStatus> feedReplies) {
+    public ReplyContext(String docId, BlockingQueue<OperationStatus> feedReplies, DocumentOperationType documentOperationType) {
         this.docId = docId;
         this.feedReplies = feedReplies;
         this.creationTime = System.currentTimeMillis();