Be able to render async error messages from event stream in json as well

2 files changed, 46 insertions, 15 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/rendering/EventRenderer.java b/container-search/src/main/java/com/yahoo/search/rendering/EventRenderer.java
index 83ae349f5a0..88a1e6c1485 100644
--- a/container-search/src/main/java/com/yahoo/search/rendering/EventRenderer.java
+++ b/container-search/src/main/java/com/yahoo/search/rendering/EventRenderer.java
@@ -64,7 +64,17 @@ public class EventRenderer extends AsynchronousSectionedRenderer<Result> {
 
     @Override
     public void data(Data data) throws IOException {
-        if (data instanceof EventStream.Event event) {
+        if (data instanceof EventStream.ErrorEvent error) {
+            generator.writeRaw("event: error\n");
+            generator.writeRaw("data: ");
+            generator.writeStartObject();
+            generator.writeStringField("source", error.source());
+            generator.writeNumberField("error", error.code());
+            generator.writeStringField("message", error.message());
+            generator.writeEndObject();
+            generator.writeRaw("\n\n");
+            generator.flush();
+        } else if (data instanceof EventStream.Event event) {
             if (RENDER_EVENT_HEADER) {
                 generator.writeRaw("event: " + event.type() + "\n");
             }
@@ -75,19 +85,6 @@ public class EventRenderer extends AsynchronousSectionedRenderer<Result> {
             generator.writeRaw("\n\n");
             generator.flush();
         }
-        else if (data instanceof ErrorHit) {
-            for (ErrorMessage error : ((ErrorHit) data).errors()) {
-                generator.writeRaw("event: error\n");
-                generator.writeRaw("data: ");
-                generator.writeStartObject();
-                generator.writeStringField("source", error.getSource());
-                generator.writeNumberField("error", error.getCode());
-                generator.writeStringField("message", error.getMessage());
-                generator.writeEndObject();
-                generator.writeRaw("\n\n");
-                generator.flush();
-            }
-        }
         // Todo: support other types of data such as search results (hits), timing and trace
     }
 
diff --git a/container-search/src/main/java/com/yahoo/search/result/EventStream.java b/container-search/src/main/java/com/yahoo/search/result/EventStream.java
index b393a91e6d0..8e6f7977d55 100644
--- a/container-search/src/main/java/com/yahoo/search/result/EventStream.java
+++ b/container-search/src/main/java/com/yahoo/search/result/EventStream.java
@@ -41,7 +41,7 @@ public class EventStream extends Hit implements DataList<Data> {
     }
 
     public void error(String source, ErrorMessage message) {
-        incoming().add(new DefaultErrorHit(source, message));
+        incoming().add(new ErrorEvent(eventCount.incrementAndGet(), source, message));
     }
 
     public void markComplete() {
@@ -117,4 +117,38 @@ public class EventStream extends Hit implements DataList<Data> {
 
     }
 
+    public static class ErrorEvent extends Event {
+
+        private final String source;
+        private final ErrorMessage message;
+
+        public ErrorEvent(int eventNumber, String source, ErrorMessage message) {
+            super(eventNumber, message.getMessage(), "error");
+            this.source = source;
+            this.message = message;
+        }
+
+        public String source() {
+            return source;
+        }
+
+        public int code() {
+            return message.getCode();
+        }
+
+        public String message() {
+            return message.getMessage();
+        }
+
+        @Override
+        public Hit asHit() {
+            Hit hit = super.asHit();
+            hit.setField("source", source);
+            hit.setField("code", message.getCode());
+            return hit;
+        }
+
+
+    }
+
 }