Introduced a new elasticsearch exception family that can hold headers

- These heades will be copied as response header on the rest response

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

@@ -19,8 +19,17 @@
 
 package org.elasticsearch;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import org.elasticsearch.common.Nullable;
+import org.elasticsearch.common.collect.Tuple;
+import org.elasticsearch.rest.HasRestHeaders;
 import org.elasticsearch.rest.RestStatus;
 
+import java.util.List;
+import java.util.Map;
+
 /**
  * A base class for all elasticsearch exceptions.
  */
@@ -152,4 +161,57 @@ public class ElasticsearchException extends RuntimeException {
             return false;
         }
     }
+
+    /**
+     * A base class for exceptions that should carry rest headers
+     */
+    @SuppressWarnings("unchecked")
+    public static class WithRestHeaders extends ElasticsearchException implements HasRestHeaders {
+
+        private final ImmutableMap<String, List<String>> headers;
+
+        public WithRestHeaders(String msg, Tuple<String, String[]>... headers) {
+            super(msg);
+            this.headers = headers(headers);
+        }
+
+        public WithRestHeaders(String msg, @Nullable ImmutableMap<String, List<String>> headers) {
+            super(msg);
+            this.headers = headers != null ? headers : ImmutableMap.<String, List<String>>of();
+        }
+
+        public WithRestHeaders(String msg, Throwable cause, Tuple<String, String[]>... headers) {
+            super(msg, cause);
+            this.headers = headers(headers);
+        }
+
+        public WithRestHeaders(String msg, Throwable cause, @Nullable ImmutableMap<String, List<String>> headers) {
+            super(msg, cause);
+            this.headers = headers != null ? headers : ImmutableMap.<String, List<String>>of();
+        }
+
+        public ImmutableMap<String, List<String>> getHeaders() {
+            return headers;
+        }
+
+        protected static Tuple<String, String[]> header(String name, String... values) {
+            return Tuple.tuple(name, values);
+        }
+
+        private static ImmutableMap<String, List<String>> headers(Tuple<String, String[]>... headers) {
+            Map<String, List<String>> map = Maps.newHashMap();
+            for (Tuple<String, String[]> header : headers) {
+                List<String> list = map.get(header.v1());
+                if (list == null) {
+                    list = Lists.newArrayList(header.v2());
+                    map.put(header.v1(), list);
+                } else {
+                    for (String value : header.v2()) {
+                        list.add(value);
+                    }
+                }
+            }
+            return ImmutableMap.copyOf(map);
+        }
+    }
 }

src/main/java/org/elasticsearch/rest/BytesRestResponse.java

@@ -93,6 +93,9 @@ public class BytesRestResponse extends RestResponse {
             this.content = builder.bytes();
             this.contentType = builder.contentType().restContentType();
         }
+        if (t instanceof HasRestHeaders) {
+            addHeaders(((HasRestHeaders) t).getHeaders());
+        }
         this.contentThreadSafe = true;
     }
 

src/main/java/org/elasticsearch/rest/HasRestHeaders.java

@@ -0,0 +1,38 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.rest;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Classes that carry rest headers should implement this interface. Specifically, exceptions that
+ * get translated to a rest response, can implement this interface and the headers will be added
+ * the the response.
+ *
+ * @see org.elasticsearch.ElasticsearchException.WithRestHeaders
+ */
+public interface HasRestHeaders {
+
+    /**
+     * @return The rest headers
+     */
+    Map<String, List<String>> getHeaders();
+}

src/main/java/org/elasticsearch/rest/RestResponse.java

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.rest;
 
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.bytes.BytesReference;
 
@@ -30,7 +32,7 @@ import java.util.Map;
 /**
  *
  */
-public abstract class RestResponse {
+public abstract class RestResponse implements HasRestHeaders {
 
     protected Map<String, List<String>> customHeaders;
 
@@ -56,6 +58,20 @@ public abstract class RestResponse {
      */
     public abstract RestStatus status();
 
+    public void addHeaders(Map<String, List<String>> headers) {
+        if (customHeaders == null) {
+            customHeaders = new HashMap<>(headers.size());
+        }
+        for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
+            List<String> values = customHeaders.get(entry.getKey());
+            if (values == null) {
+                values = Lists.newArrayList();
+                customHeaders.put(entry.getKey(), values);
+            }
+            values.addAll(entry.getValue());
+        }
+    }
+
     /**
      * Add a custom header.
      */

src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.rest;
+
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.test.ElasticsearchTestCase;
+import org.junit.Test;
+
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.notNullValue;
+
+/**
+ *
+ */
+public class BytesRestResponseTests extends ElasticsearchTestCase {
+
+    @Test
+    public void testWithHeaders() throws Exception {
+        RestRequest request = new FakeRestRequest();
+        RestChannel channel = new RestChannel(request) {
+            @Override
+            public void sendResponse(RestResponse response) {
+            }
+        };
+        BytesRestResponse response = new BytesRestResponse(channel, new ExceptionWithHeaders());
+        assertThat(response.getHeaders().get("n1"), notNullValue());
+        assertThat(response.getHeaders().get("n1"), contains("v11", "v12"));
+        assertThat(response.getHeaders().get("n2"), notNullValue());
+        assertThat(response.getHeaders().get("n2"), contains("v21", "v22"));
+    }
+
+
+    private static class ExceptionWithHeaders extends ElasticsearchException.WithRestHeaders {
+
+        ExceptionWithHeaders() {
+            super("", header("n1", "v11", "v12"), header("n2", "v21", "v22"));
+        }
+    }
+}