Carry on headers if unknow exception is a WithRestHeadersException

core/src/main/java/org/elasticsearch/ElasticsearchException.java

@@ -202,6 +202,11 @@ public class ElasticsearchException extends RuntimeException implements ToXConte
             this.headers = headers(headers);
         }
 
+        protected WithRestHeadersException(String msg, Throwable cause, Map<String, List<String>> headers) {
+            super(msg, cause);
+            this.headers = headers;
+        }
+
         public WithRestHeadersException(StreamInput in) throws IOException {
             super(in);
             int numKeys = in.readVInt();
@@ -545,6 +550,7 @@ public class ElasticsearchException extends RuntimeException implements ToXConte
                 org.elasticsearch.action.PrimaryMissingActionException.class,
                 org.elasticsearch.index.engine.CreateFailedEngineException.class,
                 org.elasticsearch.index.shard.IllegalIndexShardStateException.class,
+                WithRestHeadersException.class,
                 NotSerializableExceptionWrapper.class
         };
         Map<String, Constructor<? extends ElasticsearchException>> mapping = new HashMap<>(exceptions.length);

core/src/main/java/org/elasticsearch/common/io/stream/NotSerializableExceptionWrapper.java

@@ -20,8 +20,12 @@
 package org.elasticsearch.common.io.stream;
 
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.common.collect.Tuple;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 
 /**
  * This exception can be used to wrap a given, not serializable exception
@@ -30,12 +34,12 @@ import java.io.IOException;
  * the throwable it was created with instead of it's own. The stacktrace has no indication
  * of where this exception was created.
  */
-public final class NotSerializableExceptionWrapper extends ElasticsearchException {
+public final class NotSerializableExceptionWrapper extends ElasticsearchException.WithRestHeadersException {
 
     private final String name;
 
-    public NotSerializableExceptionWrapper(Throwable other) {
+    public NotSerializableExceptionWrapper(Throwable other, Map<String, List<String>> headers) {
-        super(other.getMessage(), other.getCause());
+        super(other.getMessage(), other.getCause(), headers);
         this.name = ElasticsearchException.getExceptionName(other);
         setStackTrace(other.getStackTrace());
         for (Throwable otherSuppressed : other.getSuppressed()) {
@@ -43,6 +47,14 @@ public final class NotSerializableExceptionWrapper extends ElasticsearchExceptio
         }
     }
 
+    public NotSerializableExceptionWrapper(WithRestHeadersException other) {
+        this(other, other.getHeaders());
+    }
+
+    public NotSerializableExceptionWrapper(Throwable other) {
+        this(other, Collections.EMPTY_MAP);
+    }
+
     public NotSerializableExceptionWrapper(StreamInput in) throws IOException {
         super(in);
         name = in.readString();

core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java

@@ -511,6 +511,9 @@ public abstract class StreamOutput extends OutputStream {
                 final String name = throwable.getClass().getName();
                 if (throwable instanceof ElasticsearchException && ElasticsearchException.isRegistered(name)) {
                     ex = (ElasticsearchException) throwable;
+                } else if (throwable instanceof ElasticsearchException.WithRestHeadersException) {
+                    // ensure we transport also the headers
+                    ex = new NotSerializableExceptionWrapper((ElasticsearchException.WithRestHeadersException)throwable);
                 } else {
                     ex = new NotSerializableExceptionWrapper(throwable);
                 }

core/src/test/java/org/elasticsearch/ExceptionSerializationTests.java

@@ -33,6 +33,7 @@ import org.elasticsearch.cluster.metadata.SnapshotId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.*;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
+import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.PathUtils;
 import org.elasticsearch.common.io.stream.*;
 import org.elasticsearch.common.transport.LocalTransportAddress;
@@ -170,6 +171,7 @@ public class ExceptionSerializationTests extends ElasticsearchTestCase {
         final Path testStartPath = PathUtils.get(ExceptionSerializationTests.class.getResource(path).toURI());
         Files.walkFileTree(testStartPath, visitor);
         assertTrue(notRegistered.remove(TestException.class));
+        assertTrue(notRegistered.remove(UnknownHeaderException.class));
         assertTrue("Classes subclassing ElasticsearchException must be registered \n" + notRegistered.toString(),
                 notRegistered.isEmpty());
         assertTrue(registered.removeAll(ElasticsearchException.getRegisteredKeys())); // check
@@ -578,4 +580,26 @@ public class ExceptionSerializationTests extends ElasticsearchTestCase {
             }
         }
     }
+
+    public void testWithRestHeadersException() throws IOException {
+        ElasticsearchException.WithRestHeadersException ex = serialize(new ElasticsearchException.WithRestHeadersException("msg", new Tuple("foo", new String[]{"foo", "bar"})));
+        assertEquals("msg", ex.getMessage());
+        assertEquals(2, ex.getHeaders().get("foo").size());
+        assertEquals("foo", ex.getHeaders().get("foo").get(0));
+        assertEquals("bar", ex.getHeaders().get("foo").get(1));
+
+        // ensure we are carrying over the headers even if not serialized
+        ElasticsearchException e = serialize((ElasticsearchException)new UnknownHeaderException("msg", new Tuple("foo", new String[]{"foo", "bar"})));
+        assertTrue(e instanceof NotSerializableExceptionWrapper);
+        assertEquals("msg", ex.getMessage());
+        assertEquals(2, ex.getHeaders().get("foo").size());
+        assertEquals("foo", ex.getHeaders().get("foo").get(0));
+        assertEquals("bar", ex.getHeaders().get("foo").get(1));
+    }
+
+    public static class UnknownHeaderException extends ElasticsearchException.WithRestHeadersException {
+        public UnknownHeaderException(String msg, Tuple<String, String[]>... headers) {
+            super(msg, headers);
+        }
+    }
 }