Fix / Data pagination background load (#619)

* Add multi batch tests to codec test suite

* Make range selector synchronous and add a test case

* Only fail data tests on memory issues inside the TRAC services

* Remove double buffering

Author: martin-traverse

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/StreamingDecoder.java

@@ -43,24 +43,12 @@ public DataPipeline.StreamApi dataInterface() {
     @Override
     public boolean isReady() {
 
-        if (context != null)
-            return context.readyToFlip() || context.readyToLoad();
-        else
-            return consumerReady();
+        return consumerReady();
     }
 
     @Override
     public void pump() {
 
-        if (context == null)
-            return;
-
-        if (context.readyToFlip())
-            context.flip();
-
-        if (context.readyToUnload() && consumerReady()) {
-            consumer().onBatch();
-            context.setUnloaded();
-        }
+        // No-op
     }
 }

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/arrow/ArrowDecoder.java

@@ -129,26 +129,17 @@ private boolean sendBatches() throws IOException {
         // Data arrives in one big buffer, so don't send it all at once
         // PUsh what the consumer requests, then wait for another call to pump()
 
-        while (context.readyToFlip() || context.readyToLoad()) {
+        while (consumerReady()) {
 
-            if (context.readyToFlip())
-                context.flip();
-
-            if (context.readyToLoad()) {
-
-                var batchAvailable = reader.loadNextBatch();
-
-                if (!batchAvailable)
-                    return true;
+            var batchAvailable = reader.loadNextBatch();
 
+            if (batchAvailable) {
                 context.setLoaded();
-
-                if (context.readyToFlip())
-                    context.flip();
-            }
-
-            if (context.readyToUnload() && consumerReady())
                 consumer().onBatch();
+            }
+            else {
+                return true;
+            }
         }
 
         return false;

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/arrow/ArrowFileEncoder.java

@@ -34,6 +34,6 @@ public ArrowFileEncoder(BufferAllocator allocator) {
     @Override
     protected ArrowWriter createWriter(ArrowVsrContext context, BufferAllocator allocator) {
         var out = new ByteOutputChannel(context.getAllocator(), consumer()::onNext);
-        return new ArrowFileWriterExt(context.getFrontBuffer(), context.getDictionaries(), out);
+        return new ArrowFileWriterExt(context.getVsr(), context.getDictionaries(), out);
     }
 }

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/arrow/ArrowStreamDecoder.java

@@ -124,7 +124,7 @@ public void pump() {
 
     private void sendBatches() throws IOException {
 
-        if (isDone())
+        if (!consumerReady() || isDone())
             return;
 
         if (context == null) {
@@ -143,37 +143,22 @@ else if (messageReader.hasEos()) {
             }
         }
 
-        if (context.readyToFlip())
-            context.flip();
-
-        if (context.readyToLoad()) {
-
-            if (messageReader.hasMessage(MessageHeader.RecordBatch)) {
-
-                arrowReader.loadNextBatch();
-                context.setLoaded();
-
-                if (context.readyToFlip())
-                    context.flip();
-            }
-            else if (messageReader.hasEos()) {
+        while (consumerReady() && messageReader.hasMessage(MessageHeader.RecordBatch)) {
+            arrowReader.loadNextBatch();
+            consumer().onBatch();
+        }
 
-                if (messageReader.hasMessage()) {
-                    var error = new EDataCorruption("Arrow decoding failed, unexpected messages after EOS marker");
-                    onError(error);
-                    return;
-                }
+        if (consumerReady() && messageReader.hasEos()) {
 
-                if (!isDone())
-                    markAsDone();
+            if (messageReader.hasMessage()) {
+                var error = new EDataCorruption("Arrow decoding failed, unexpected messages after EOS marker");
+                onError(error);
+                return;
             }
-        }
 
-        if (context.readyToUnload() && consumerReady())
-            consumer().onBatch();
-
-        if (isDone() && !context.readyToUnload())
+            markAsDone();
             consumer().onComplete();
+        }
     }
 
     @Override

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/arrow/ArrowStreamEncoder.java

@@ -34,6 +34,6 @@ public ArrowStreamEncoder(BufferAllocator allocator) {
     @Override
     protected ArrowWriter createWriter(ArrowVsrContext context, BufferAllocator allocator) {
         var out = new ByteOutputChannel(context.getAllocator(), consumer()::onNext);
-        return new ArrowStreamWriter(context.getFrontBuffer(), context.getDictionaries(), out);
+        return new ArrowStreamWriter(context.getVsr(), context.getDictionaries(), out);
     }
 }

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/text/BaseTextDecoder.java

@@ -205,11 +205,9 @@ boolean doParse() throws IOException {
         while (reader.readBatch()) {
 
             context.setLoaded();
-            context.flip();
             consumer().onBatch();
-            context.setUnloaded();
 
-            reader.resetBatch(context.getBackBuffer());
+            reader.resetBatch(context.getVsr());
         }
 
         return reader.endOfStream();

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/text/BaseTextEncoder.java

@@ -73,7 +73,7 @@ public void onStart(ArrowVsrContext context) {
             this.out = new ByteOutputStream(allocator, consumer()::onNext);
 
             this.writer = new TextFileWriter(
-                    context.getFrontBuffer(),
+                    context.getVsr(),
                     context.getDictionaries(),
                     this.out,
                     configForSchema(this.config, context.getSchema()));
@@ -112,7 +112,7 @@ public void onBatch() {
             if (log.isTraceEnabled())
                 log.trace("JSON ENCODER: onNext()");
 
-            var batch = context.getFrontBuffer();
+            var batch = context.getVsr();
 
             writer.resetBatch(batch);
             writer.writeBatch();

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/codec/text/BufferedTextDecoder.java

@@ -172,11 +172,9 @@ boolean doParse() throws Exception {
         while (consumerReady() && reader.readBatch()) {
 
             context.setLoaded();
-            context.flip();
             consumer().onBatch();
-            context.setUnloaded();
 
-            reader.resetBatch(context.getBackBuffer());
+            reader.resetBatch(context.getVsr());
         }
 
         return reader.endOfStream();

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/data/ArrowVsrContext.java

@@ -30,15 +30,12 @@ public class ArrowVsrContext {
     private final ArrowVsrSchema schema;
 
     private final BufferAllocator allocator;
-    private final VectorSchemaRoot front;
-    private final VectorSchemaRoot back;
+    private final VectorSchemaRoot vsr;
     private final DictionaryProvider dictionaries;
 
     private final boolean ownership;
 
-    private boolean backBufferLoaded;
-    private boolean frontBufferAvailable;
-    private boolean frontBufferUnloaded;
+    private boolean loaded;
 
     public static ArrowVsrContext forSource(VectorSchemaRoot source, DictionaryProvider dictionaries, BufferAllocator allocator) {
 
@@ -56,8 +53,7 @@ private ArrowVsrContext(VectorSchemaRoot source, DictionaryProvider dictionaries
         this.schema = new ArrowVsrSchema(source.getSchema(), dictionaries);
 
         this.allocator = allocator;
-        this.back = source;
-        this.front = back;  // No double buffering yet
+        this.vsr = source;
         this.dictionaries = dictionaries;
 
         this.ownership = takeOwnership;
@@ -84,8 +80,7 @@ private ArrowVsrContext(ArrowVsrSchema schema, BufferAllocator allocator) {
         var root = new VectorSchemaRoot(fields, vectors);
         root.allocateNew();
 
-        this.back = root;
-        this.front = root;  // No double buffering yet
+        this.vsr = root;
 
         // Use pre-defined dictionaries from the schema (if there are any)
         this.dictionaries = schema.dictionaries();
@@ -102,55 +97,39 @@ public BufferAllocator getAllocator() {
         return allocator;
     }
 
-    public VectorSchemaRoot getBackBuffer() {
-        return back;
-    }
-
-    public VectorSchemaRoot getFrontBuffer() {
-        return front;
+    public VectorSchemaRoot getVsr() {
+        return vsr;
     }
 
     public DictionaryProvider getDictionaries() {
         return dictionaries;
     }
-
-    public boolean readyToLoad() {
-        return ! backBufferLoaded;
-    }
-
     public void setRowCount(int nRows) {
 
-        back.setRowCount(nRows);
+        vsr.setRowCount(nRows);
     }
 
     public void setLoaded() {
-        backBufferLoaded = true;
+        loaded = true;
     }
 
-    public boolean readyToFlip() {
-        return backBufferLoaded && (frontBufferUnloaded || !frontBufferAvailable);
+    public void setUnloaded() {
+        loaded = false;
     }
 
-    public void flip() {
-        frontBufferAvailable = true;
-        frontBufferUnloaded = false;
+    public boolean readyToLoad() {
+        return !loaded;
     }
 
     public boolean readyToUnload() {
-        return frontBufferAvailable && !frontBufferUnloaded;
-    }
-
-    public void setUnloaded() {
-        frontBufferUnloaded = true;
-        frontBufferAvailable = false;
-        backBufferLoaded = false;
+        return loaded;
     }
 
     public void close() {
 
         if (ownership) {
 
-            back.close();
+            vsr.close();
 
             if (dictionaries != null) {
                 for (var dictionaryId : dictionaries.getDictionaryIds()) {

tracdap-libs/tracdap-lib-data/src/main/java/org/finos/tracdap/common/data/pipeline/CounterStage.java

@@ -53,7 +53,7 @@ public long getRowCount() {
     @Override
     public void onBatch() {
         batchCount++;
-        rowCount += batch.getFrontBuffer().getRowCount();
+        rowCount += batch.getVsr().getRowCount();
         consumer().onBatch();
     }