ESQL: Read from the BlockFactory

test/framework/src/main/java/org/elasticsearch/test/AbstractSerializationTestCase.java

@@ -39,6 +39,7 @@ public final void testFromXContent() throws IOException {
             .randomFieldsExcludeFilter(getRandomFieldsExcludeFilter())
             .assertEqualsConsumer(this::assertEqualInstances)
             .assertToXContentEquivalence(assertToXContentEquivalence())
+            .dispose(this::dispose)
             .test();
     }
 
@@ -61,41 +62,45 @@ public final void testConcurrentToXContent() throws IOException, InterruptedExce
             () -> randomFrom(XContentType.values())
         );
         T testInstance = createXContextTestInstance(xContentType);
-        ToXContent.Params params = new ToXContent.DelegatingMapParams(
-            singletonMap(RestSearchAction.TYPED_KEYS_PARAM, "true"),
-            getToXContentParams()
-        );
-        boolean humanReadable = randomBoolean();
-        BytesRef firstTimeBytes = toXContent(asXContent(testInstance), xContentType, params, humanReadable).toBytesRef();
-
-        /*
-         * 500 rounds seems to consistently reproduce the issue on Nik's
-         * laptop. Larger numbers are going to be slower but more likely
-         * to reproduce the issue.
-         */
-        int rounds = scaledRandomIntBetween(300, 5000);
-        concurrentTest(() -> {
-            try {
-                for (int r = 0; r < rounds; r++) {
-                    BytesRef thisRoundBytes = toXContent(asXContent(testInstance), xContentType, params, humanReadable).toBytesRef();
-                    if (firstTimeBytes.bytesEquals(thisRoundBytes)) {
-                        continue;
-                    }
-                    StringBuilder error = new StringBuilder("Failed to round trip over ");
-                    if (humanReadable) {
-                        error.append("human readable ");
+        try {
+            ToXContent.Params params = new ToXContent.DelegatingMapParams(
+                singletonMap(RestSearchAction.TYPED_KEYS_PARAM, "true"),
+                getToXContentParams()
+            );
+            boolean humanReadable = randomBoolean();
+            BytesRef firstTimeBytes = toXContent(asXContent(testInstance), xContentType, params, humanReadable).toBytesRef();
+
+            /*
+             * 500 rounds seems to consistently reproduce the issue on Nik's
+             * laptop. Larger numbers are going to be slower but more likely
+             * to reproduce the issue.
+             */
+            int rounds = scaledRandomIntBetween(300, 5000);
+            concurrentTest(() -> {
+                try {
+                    for (int r = 0; r < rounds; r++) {
+                        BytesRef thisRoundBytes = toXContent(asXContent(testInstance), xContentType, params, humanReadable).toBytesRef();
+                        if (firstTimeBytes.bytesEquals(thisRoundBytes)) {
+                            continue;
+                        }
+                        StringBuilder error = new StringBuilder("Failed to round trip over ");
+                        if (humanReadable) {
+                            error.append("human readable ");
+                        }
+                        error.append(xContentType);
+                        error.append("\nCanonical is:\n").append(Strings.toString(asXContent(testInstance), true, true));
+                        boolean showBytes = xContentType.xContent() == CborXContent.cborXContent;
+                        error.append("\nWanted : ").append(showBytes ? firstTimeBytes : firstTimeBytes.utf8ToString());
+                        error.append("\nBut got: ").append(showBytes ? thisRoundBytes : thisRoundBytes.utf8ToString());
+                        fail(error.toString());
                     }
-                    error.append(xContentType);
-                    error.append("\nCanonical is:\n").append(Strings.toString(asXContent(testInstance), true, true));
-                    boolean showBytes = xContentType.xContent() == CborXContent.cborXContent;
-                    error.append("\nWanted : ").append(showBytes ? firstTimeBytes : firstTimeBytes.utf8ToString());
-                    error.append("\nBut got: ").append(showBytes ? thisRoundBytes : thisRoundBytes.utf8ToString());
-                    fail(error.toString());
+                } catch (IOException e) {
+                    throw new AssertionError(e);
                 }
-            } catch (IOException e) {
-                throw new AssertionError(e);
-            }
-        });
+            });
+        } finally {
+            dispose(testInstance);
+        }
     }
 
     protected abstract ToXContent asXContent(T instance);

test/framework/src/main/java/org/elasticsearch/test/AbstractWireTestCase.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
@@ -54,10 +55,20 @@ public abstract class AbstractWireTestCase<T> extends ESTestCase {
      */
     public final void testEqualsAndHashcode() {
         for (int runs = 0; runs < NUMBER_OF_TEST_RUNS; runs++) {
-            EqualsHashCodeTestUtils.checkEqualsAndHashCode(createTestInstance(), this::copyInstance, this::mutateInstance);
+            T testInstance = createTestInstance();
+            try {
+                EqualsHashCodeTestUtils.checkEqualsAndHashCode(testInstance, this::copyInstance, this::mutateInstance, this::dispose);
+            } finally {
+                dispose(testInstance);
+            }
         }
     }
 
+    /**
+     * Dispose of the copy, usually {@link Releasable#close} or a noop.
+     */
+    protected void dispose(T t) {}
+
     /**
      * Calls {@link Object#equals} on equal objects on many threads and verifies
      * they all return true. Folks tend to assume this is true about
@@ -67,19 +78,27 @@ public final void testEqualsAndHashcode() {
      */
     public final void testConcurrentEquals() throws IOException, InterruptedException, ExecutionException {
         T testInstance = createTestInstance();
-        T copy = copyInstance(testInstance);
-
-        /*
-         * 500 rounds seems to consistently reproduce the issue on Nik's
-         * laptop. Larger numbers are going to be slower but more likely
-         * to reproduce the issue.
-         */
-        int rounds = scaledRandomIntBetween(300, 5000);
-        concurrentTest(() -> {
-            for (int r = 0; r < rounds; r++) {
-                assertEquals(testInstance, copy);
+        try {
+            T copy = copyInstance(testInstance);
+            try {
+
+                /*
+                 * 500 rounds seems to consistently reproduce the issue on Nik's
+                 * laptop. Larger numbers are going to be slower but more likely
+                 * to reproduce the issue.
+                 */
+                int rounds = scaledRandomIntBetween(300, 5000);
+                concurrentTest(() -> {
+                    for (int r = 0; r < rounds; r++) {
+                        assertEquals(testInstance, copy);
+                    }
+                });
+            } finally {
+                dispose(copy);
             }
-        });
+        } finally {
+            dispose(testInstance);
+        }
     }
 
     /**
@@ -111,25 +130,34 @@ protected void concurrentTest(Runnable r) throws InterruptedException, Execution
      */
     public final void testConcurrentHashCode() throws InterruptedException, ExecutionException {
         T testInstance = createTestInstance();
-        int firstHashCode = testInstance.hashCode();
-
-        /*
-         * 500 rounds seems to consistently reproduce the issue on Nik's
-         * laptop. Larger numbers are going to be slower but more likely
-         * to reproduce the issue.
-         */
-        int rounds = scaledRandomIntBetween(300, 5000);
-        concurrentTest(() -> {
-            for (int r = 0; r < rounds; r++) {
-                assertEquals(firstHashCode, testInstance.hashCode());
-            }
-        });
+        try {
+            int firstHashCode = testInstance.hashCode();
+
+            /*
+             * 500 rounds seems to consistently reproduce the issue on Nik's
+             * laptop. Larger numbers are going to be slower but more likely
+             * to reproduce the issue.
+             */
+            int rounds = scaledRandomIntBetween(300, 5000);
+            concurrentTest(() -> {
+                for (int r = 0; r < rounds; r++) {
+                    assertEquals(firstHashCode, testInstance.hashCode());
+                }
+            });
+        } finally {
+            dispose(testInstance);
+        }
     }
 
     public void testToString() throws Exception {
-        final String toString = createTestInstance().toString();
-        assertNotNull(toString);
-        assertThat(toString, not(emptyString()));
+        T testInstance = createTestInstance();
+        try {
+            final String toString = testInstance.toString();
+            assertNotNull(toString);
+            assertThat(toString, not(emptyString()));
+        } finally {
+            dispose(testInstance);
+        }
     }
 
     /**
@@ -138,7 +166,11 @@ public void testToString() throws Exception {
     public final void testSerialization() throws IOException {
         for (int runs = 0; runs < NUMBER_OF_TEST_RUNS; runs++) {
             T testInstance = createTestInstance();
-            assertSerialization(testInstance);
+            try {
+                assertSerialization(testInstance);
+            } finally {
+                dispose(testInstance);
+            }
         }
     }
 
@@ -155,22 +187,25 @@ public final void testSerialization() throws IOException {
      */
     public final void testConcurrentSerialization() throws InterruptedException, ExecutionException {
         T testInstance = createTestInstance();
-
-        /*
-         * 500 rounds seems to consistently reproduce the issue on Nik's
-         * laptop. Larger numbers are going to be slower but more likely
-         * to reproduce the issue.
-         */
-        int rounds = scaledRandomIntBetween(300, 2000);
-        concurrentTest(() -> {
-            try {
-                for (int r = 0; r < rounds; r++) {
-                    assertSerialization(testInstance);
+        try {
+            /*
+             * 500 rounds seems to consistently reproduce the issue on Nik's
+             * laptop. Larger numbers are going to be slower but more likely
+             * to reproduce the issue.
+             */
+            int rounds = scaledRandomIntBetween(300, 2000);
+            concurrentTest(() -> {
+                try {
+                    for (int r = 0; r < rounds; r++) {
+                        assertSerialization(testInstance);
+                    }
+                } catch (IOException e) {
+                    throw new AssertionError("error serializing", e);
                 }
-            } catch (IOException e) {
-                throw new AssertionError("error serializing", e);
-            }
-        });
+            });
+        } finally {
+            dispose(testInstance);
+        }
     }
 
     /**
@@ -187,7 +222,11 @@ protected final void assertSerialization(T testInstance) throws IOException {
      */
     protected final void assertSerialization(T testInstance, TransportVersion version) throws IOException {
         T deserializedInstance = copyInstance(testInstance, version);
-        assertEqualInstances(testInstance, deserializedInstance);
+        try {
+            assertEqualInstances(testInstance, deserializedInstance);
+        } finally {
+            dispose(deserializedInstance);
+        }
     }
 
     /**

test/framework/src/main/java/org/elasticsearch/test/AbstractXContentTestCase.java

@@ -24,6 +24,7 @@
 
 import java.io.IOException;
 import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
@@ -125,6 +126,7 @@ public static class XContentTester<T> {
             assertEquals(expectedInstance.hashCode(), newInstance.hashCode());
         };
         private boolean assertToXContentEquivalence = true;
+        private Consumer<T> dispose = t -> {};
 
         private XContentTester(
             CheckedBiFunction<XContent, BytesReference, XContentParser, IOException> createParser,
@@ -142,24 +144,32 @@ public void test() throws IOException {
             for (int runs = 0; runs < numberOfTestRuns; runs++) {
                 XContentType xContentType = randomFrom(XContentType.values()).canonical();
                 T testInstance = instanceSupplier.apply(xContentType);
-                BytesReference originalXContent = toXContent.apply(testInstance, xContentType);
-                BytesReference shuffledContent = insertRandomFieldsAndShuffle(
-                    originalXContent,
-                    xContentType,
-                    supportsUnknownFields,
-                    shuffleFieldsExceptions,
-                    randomFieldsExcludeFilter,
-                    createParser
-                );
-                XContentParser parser = createParser.apply(XContentFactory.xContent(xContentType), shuffledContent);
-                T parsed = fromXContent.apply(parser);
-                assertEqualsConsumer.accept(testInstance, parsed);
-                if (assertToXContentEquivalence) {
-                    assertToXContentEquivalent(
-                        toXContent.apply(testInstance, xContentType),
-                        toXContent.apply(parsed, xContentType),
-                        xContentType
+                try {
+                    BytesReference originalXContent = toXContent.apply(testInstance, xContentType);
+                    BytesReference shuffledContent = insertRandomFieldsAndShuffle(
+                        originalXContent,
+                        xContentType,
+                        supportsUnknownFields,
+                        shuffleFieldsExceptions,
+                        randomFieldsExcludeFilter,
+                        createParser
                     );
+                    XContentParser parser = createParser.apply(XContentFactory.xContent(xContentType), shuffledContent);
+                    T parsed = fromXContent.apply(parser);
+                    try {
+                        assertEqualsConsumer.accept(testInstance, parsed);
+                        if (assertToXContentEquivalence) {
+                            assertToXContentEquivalent(
+                                toXContent.apply(testInstance, xContentType),
+                                toXContent.apply(parsed, xContentType),
+                                xContentType
+                            );
+                        }
+                    } finally {
+                        dispose.accept(parsed);
+                    }
+                } finally {
+                    dispose.accept(testInstance);
                 }
             }
         }
@@ -193,6 +203,11 @@ public XContentTester<T> assertToXContentEquivalence(boolean assertToXContentEqu
             this.assertToXContentEquivalence = assertToXContentEquivalence;
             return this;
         }
+
+        public XContentTester<T> dispose(Consumer<T> dispose) {
+            this.dispose = dispose;
+            return this;
+        }
     }
 
     public static <T extends ToXContent> void testFromXContent(