All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.deephaven.stream.TableStreamPublisherImpl Maven / Gradle / Ivy

There is a newer version: 0.37.1
Show newest version
package io.deephaven.stream;

import io.deephaven.chunk.WritableChunk;
import io.deephaven.chunk.attributes.Values;
import io.deephaven.engine.context.ExecutionContext;
import io.deephaven.engine.rowset.RowSequence;
import io.deephaven.engine.rowset.RowSequence.Iterator;
import io.deephaven.engine.rowset.RowSet;
import io.deephaven.engine.table.ChunkSource.FillContext;
import io.deephaven.engine.table.ColumnDefinition;
import io.deephaven.engine.table.ColumnSource;
import io.deephaven.engine.table.SharedContext;
import io.deephaven.engine.table.Table;
import io.deephaven.engine.table.TableDefinition;
import io.deephaven.engine.table.impl.NotificationStepSource;
import io.deephaven.engine.table.impl.remote.ConstructSnapshot;
import io.deephaven.engine.table.impl.remote.ConstructSnapshot.SnapshotFunction;
import io.deephaven.engine.table.impl.remote.ConstructSnapshot.State;
import io.deephaven.engine.table.impl.sources.ReinterpretUtils;
import io.deephaven.engine.util.input.InputTableStatusListener;
import io.deephaven.engine.util.input.InputTableUpdater;
import io.deephaven.util.SafeCloseable;
import io.deephaven.util.SafeCloseableArray;
import org.jetbrains.annotations.NotNull;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.stream.Stream;

class TableStreamPublisherImpl implements StreamPublisher {

    private final String name;
    private final TableDefinition definition;
    private final Runnable onFlushCallback;
    private final Runnable onShutdownCallback;
    private final int chunkSize;
    private StreamConsumer consumer;

    TableStreamPublisherImpl(String name, TableDefinition definition, Runnable onFlushCallback,
            Runnable onShutdownCallback, int chunkSize) {
        if (chunkSize <= 0) {
            throw new IllegalArgumentException("chunkSize must be positive");
        }
        this.name = Objects.requireNonNull(name);
        this.definition = Objects.requireNonNull(definition);
        this.onFlushCallback = onFlushCallback;
        this.onShutdownCallback = onShutdownCallback;
        this.chunkSize = chunkSize;
    }

    @Override
    public void register(@NotNull StreamConsumer consumer) {
        if (this.consumer != null) {
            throw new IllegalStateException("Can not register multiple StreamConsumers.");
        }
        this.consumer = Objects.requireNonNull(consumer);
    }

    public TableDefinition definition() {
        return definition;
    }

    public void add(Table table) {
        // Note: we expect an exception in FillChunks / io.deephaven.engine.table.Table#getColumnSource if the table is
        // incompatible instead of a heavier-weight
        // io.deephaven.engine.table.TableDefinition checkCompatibility / checkMutualCompatibility
        // we could do here.
        final FillChunks fillChunks = new FillChunks(table);
        try (final SafeCloseable ignored =
                ExecutionContext.getContext().withUpdateGraph(table.getUpdateGraph()).open()) {
            ConstructSnapshot.callDataSnapshotFunction(
                    TableStreamPublisherImpl.class.getSimpleName() + "-" + name,
                    ConstructSnapshot.makeSnapshotControl(false, table.isRefreshing(), (NotificationStepSource) table),
                    fillChunks);
        }
        consumer.accept(fillChunks.outstandingChunks);
    }

    public void publishFailure(Throwable e) {
        consumer.acceptFailure(e);
    }

    @Override
    public void flush() {
        // even though this impl itself doesn't hold onto any data, the publisher may be batching and want to call #add
        if (onFlushCallback != null) {
            onFlushCallback.run();
        }
    }

    @Override
    public void shutdown() {
        if (onShutdownCallback != null) {
            onShutdownCallback.run();
        }
    }

    public InputTableUpdater inputTableUpdater() {
        return new InputTableAdapter();
    }

    private class InputTableAdapter implements InputTableUpdater {
        @Override
        public TableDefinition getTableDefinition() {
            return definition;
        }

        @Override
        public void add(Table newData) {
            TableStreamPublisherImpl.this.add(newData);
        }

        @Override
        public void addAsync(Table newData, InputTableStatusListener listener) {
            try {
                TableStreamPublisherImpl.this.add(newData);
            } catch (Throwable t) {
                listener.onError(t);
                return;
            }
            listener.onSuccess();
        }

        @Override
        public List getKeyNames() {
            return Collections.emptyList();
        }

    }

    private class FillChunks implements SnapshotFunction {
        private final Table table;
        private final ColumnSource[] sources;
        private final List[]> outstandingChunks;

        public FillChunks(Table table) {
            this.table = Objects.requireNonNull(table);
            this.sources = new ColumnSource[definition.numColumns()];
            // sources is in the same order as definition columns
            int i = 0;
            for (ColumnDefinition column : definition.getColumns()) {
                sources[i++] = ReinterpretUtils.maybeConvertToPrimitive(
                        table.getColumnSource(column.getName(), column.getDataType(), column.getComponentType()));
            }
            this.outstandingChunks = new ArrayList<>();
        }

        @Override
        public boolean call(boolean usePrev, long beforeClockValue) {
            final State state = ConstructSnapshot.state();
            // We are calling reset() before returning false, even though it is not technically necessary with our
            // call to reset() here. That said, the reset() before returning false saves a little bit of time during the
            // next snapshot attempt which improves our chances of getting a consistent snapshot.
            //
            // Even though we are calling reset() before we return false, we still need to check it here as it's
            // possible this function returned true, but the construct snapshot implementation re-invokes us because it
            // was an inconsistent snapshot.
            reset();
            final RowSet rowSet = usePrev
                    ? table.getRowSet().prev()
                    : table.getRowSet();
            final long initialSize = rowSet.size();
            final int numColumns = sources.length;
            final FillContext[] fillContexts = new FillContext[numColumns];
            try (
                    final SharedContext sharedContext = numColumns > 1 ? SharedContext.makeSharedContext() : null;
                    final SafeCloseable ignored = new SafeCloseableArray<>(fillContexts);
                    final Iterator rows = rowSet.getRowSequenceIterator()) {
                {
                    final int fillContextSize = (int) Math.min(chunkSize, initialSize);
                    for (int i = 0; i < numColumns; i++) {
                        fillContexts[i] = sources[i].makeFillContext(fillContextSize, sharedContext);
                    }
                }
                long remaining = initialSize;
                while (rows.hasMore()) {
                    if (sharedContext != null) {
                        sharedContext.reset();
                    }
                    assertTrue(state, remaining > 0, "remaining > 0");
                    final RowSequence rowSeq = rows.getNextRowSequenceWithLength(chunkSize);
                    final int rowSeqSize = rowSeq.intSize();
                    assertTrue(state, rowSeqSize == Math.min(chunkSize, remaining),
                            "rowSeqSize == Math.min(chunkSize, remaining)");
                    remaining -= rowSeqSize;
                    final WritableChunk[] sinks =
                            StreamChunkUtils.makeChunksForDefinition(definition, rowSeqSize);
                    // Note: adding to list ASAP to ensure they get cleaned up via #reset if there is an exception
                    // during filling.
                    outstandingChunks.add(sinks);
                    for (int i = 0; i < numColumns; ++i) {
                        if (usePrev) {
                            sources[i].fillPrevChunk(fillContexts[i], sinks[i], rowSeq);
                        } else {
                            sources[i].fillChunk(fillContexts[i], sinks[i], rowSeq);
                        }
                        if (state.concurrentAttemptInconsistent()) {
                            reset();
                            return false;
                        }
                    }
                }
                assertTrue(state, remaining == 0, "remaining == 0");
            }
            return true;
        }

        private void assertTrue(State state, boolean condition, String message) {
            if (!condition) {
                reset();
                // We are really hoping this throws an error.
                state.failIfConcurrentAttemptInconsistent();
                // This is bad.
                throw new Error("Found broken assertion not due to inconsistent attempt: " + message);
            }
        }

        private void reset() {
            SafeCloseable.closeAll(outstandingChunks.stream().flatMap(Stream::of));
            outstandingChunks.clear();
        }
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy