Improve robustness of Query Result serializations (elastic#54692) (elastic#55028)

Makes query result serialization more robust by propagating possible
IOExceptions that can occur during shard level result serialization to the
caller instead of throwing AssertionError that is not intercepted.

Fixes elastic#54665

Author: imotov

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

@@ -808,7 +808,7 @@ public void writeGenericValue(@Nullable Object value) throws IOException {
         if (writer != null) {
             writer.write(this, value);
         } else {
-            throw new IOException("can not write type [" + type + "]");
+            throw new IllegalArgumentException("can not write type [" + type + "]");
         }
     }
 

server/src/main/java/org/elasticsearch/indices/IndicesRequestCache.java

@@ -28,6 +28,7 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.RamUsageEstimator;
+import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.cache.Cache;
 import org.elasticsearch.common.cache.CacheBuilder;
@@ -43,6 +44,7 @@
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 
 import java.io.Closeable;
+import java.io.IOException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
@@ -115,8 +117,8 @@ public void onRemoval(RemovalNotification<Key, BytesReference> notification) {
     // NORELEASE The cacheKeyRenderer has been added in order to debug
     // https://github.com/elastic/elasticsearch/issues/32827, it should be
     // removed when this issue is solved
-    BytesReference getOrCompute(CacheEntity cacheEntity, Supplier<BytesReference> loader,
-            DirectoryReader reader, BytesReference cacheKey, Supplier<String> cacheKeyRenderer) throws Exception {
+    BytesReference getOrCompute(CacheEntity cacheEntity, CheckedSupplier<BytesReference, IOException> loader,
+                                DirectoryReader reader, BytesReference cacheKey, Supplier<String> cacheKeyRenderer) throws Exception {
         assert reader.getReaderCacheHelper() != null;
         final Key key =  new Key(cacheEntity, reader.getReaderCacheHelper().getKey(), cacheKey);
         Loader cacheLoader = new Loader(cacheEntity, loader);
@@ -157,10 +159,10 @@ void invalidate(CacheEntity cacheEntity, DirectoryReader reader, BytesReference
     private static class Loader implements CacheLoader<Key, BytesReference> {
 
         private final CacheEntity entity;
-        private final Supplier<BytesReference> loader;
+        private final CheckedSupplier<BytesReference, IOException> loader;
         private boolean loaded;
 
-        Loader(CacheEntity entity, Supplier<BytesReference> loader) {
+        Loader(CacheEntity entity, CheckedSupplier<BytesReference, IOException> loader) {
             this.entity = entity;
             this.loader = loader;
         }

server/src/main/java/org/elasticsearch/indices/IndicesService.java

@@ -45,7 +45,9 @@
 import org.elasticsearch.cluster.routing.RecoverySource;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.CheckedConsumer;
 import org.elasticsearch.common.CheckedFunction;
+import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.bytes.BytesReference;
@@ -1376,12 +1378,7 @@ public void loadIntoContext(ShardSearchRequest request, SearchContext context, Q
             () -> "Shard: " + request.shardId() + "\nSource:\n" + request.source(),
             out -> {
             queryPhase.execute(context);
-            try {
-                context.queryResult().writeToNoId(out);
-
-            } catch (IOException e) {
-                throw new AssertionError("Could not serialize response", e);
-            }
+            context.queryResult().writeToNoId(out);
             loadedFromCache[0] = false;
         });
 
@@ -1420,9 +1417,9 @@ public ByteSizeValue getTotalIndexingBufferBytes() {
      * @return the contents of the cache or the result of calling the loader
      */
     private BytesReference cacheShardLevelResult(IndexShard shard, DirectoryReader reader, BytesReference cacheKey,
-            Supplier<String> cacheKeyRenderer, Consumer<StreamOutput> loader) throws Exception {
+            Supplier<String> cacheKeyRenderer, CheckedConsumer<StreamOutput, IOException> loader) throws Exception {
         IndexShardCacheEntity cacheEntity = new IndexShardCacheEntity(shard);
-        Supplier<BytesReference> supplier = () -> {
+        CheckedSupplier<BytesReference, IOException> supplier = () -> {
             /* BytesStreamOutput allows to pass the expected size but by default uses
              * BigArrays.PAGE_SIZE_IN_BYTES which is 16k. A common cached result ie.
              * a date histogram with 3 buckets is ~100byte so 16k might be very wasteful

server/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java

@@ -349,6 +349,13 @@ public void testSimpleStreams() throws Exception {
         assertThat(jdt.getZonedDateTime().toInstant().toEpochMilli(), equalTo(123456L));
         assertThat(jdt.getZonedDateTime().getZone(), equalTo(ZoneId.of("America/Los_Angeles")));
         assertEquals(0, in.available());
+        IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> out.writeGenericValue(new Object() {
+            @Override
+            public String toString() {
+                return "This object cannot be serialized by writeGeneric method";
+            }
+        }));
+        assertThat(ex.getMessage(), containsString("can not write type"));
         in.close();
         out.close();
     }

server/src/test/java/org/elasticsearch/indices/IndicesRequestCacheTests.java

@@ -32,6 +32,7 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.common.bytes.AbstractBytesReference;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
@@ -49,7 +50,6 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Supplier;
 
 public class IndicesRequestCacheTests extends ESTestCase {
 
@@ -331,7 +331,7 @@ public Iterable<Field> newDoc(int id, String value) {
             StringField.TYPE_STORED));
     }
 
-    private static class Loader implements Supplier<BytesReference> {
+    private static class Loader implements CheckedSupplier<BytesReference, IOException> {
 
         private final DirectoryReader reader;
         private final int id;