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

com.kolibrifx.plovercrest.server.internal.protocol.StreamConnection Maven / Gradle / Ivy

The newest version!
/*
 * Copyright (c) 2010-2017, KolibriFX AS. Licensed under the Apache License, version 2.0.
 */

package com.kolibrifx.plovercrest.server.internal.protocol;

import static com.kolibrifx.plovercrest.client.remote.DefaultPlovercrestRemote.PROTOCOL_VERSION;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousCloseException;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.log4j.Logger;
import com.kolibrifx.plovercrest.client.MultiTableWriter.WriterMode;
import com.kolibrifx.plovercrest.client.PlovercrestAccessException;
import com.kolibrifx.plovercrest.client.PlovercrestException;
import com.kolibrifx.plovercrest.client.RangeQuery;
import com.kolibrifx.plovercrest.client.RangeQuery.RangeType;
import com.kolibrifx.plovercrest.client.TableMetadata;
import com.kolibrifx.plovercrest.client.internal.ReaderStrategy;
import com.kolibrifx.plovercrest.client.internal.RemoteCommand;
import com.kolibrifx.plovercrest.client.internal.SubscriptionQuery;
import com.kolibrifx.plovercrest.client.internal.SubscriptionQuery.QueryKind;
import com.kolibrifx.plovercrest.client.internal.protocol.GrowableByteBuffer;
import com.kolibrifx.plovercrest.client.internal.protocol.IndexedError;
import com.kolibrifx.plovercrest.client.internal.remote.RemoteResponse;
import com.kolibrifx.plovercrest.client.internal.shared.PlovercrestConstants;
import com.kolibrifx.plovercrest.client.internal.shared.TableMetadataUtils;
import com.kolibrifx.plovercrest.client.remote.DefaultPlovercrestRemote;
import com.kolibrifx.plovercrest.server.TableInfo;
import com.kolibrifx.plovercrest.server.internal.EngineAdapter;
import com.kolibrifx.plovercrest.server.internal.EngineAdapter.ChunkedDataHandler;
import com.kolibrifx.plovercrest.server.internal.EngineAdapter.ReadRequest;
import com.kolibrifx.plovercrest.server.internal.ReaderPositionUtils;
import com.kolibrifx.plovercrest.server.internal.protocol.ThreadedSubscriberEntryWriter.SubscriberWriteHandler;
import com.kolibrifx.plovercrest.server.security.AccessControlFilter;
import com.kolibrifx.plovercrest.server.security.AccessKind;
import com.kolibrifx.plovercrest.server.security.ClientInfo;
import com.kolibrifx.plovercrest.server.streams.ReaderPosition;
import com.kolibrifx.plovercrest.server.streams.ReaderPosition.PositionType;
import com.kolibrifx.plovercrest.server.streams.StreamInfo;

/**
 * Server endpoint for a plovercrest stream protocol, see
 * {@link com.kolibrifx.plovercrest.client.internal.protocol.AbstractStreamClient} for details.
 */
public class StreamConnection extends Thread {
    private static final Logger log = Logger.getLogger(StreamConnection.class);
    private static final int ACK = Integer.MIN_VALUE;
    private static final long READ_TIMEOUT_MILLISECONDS = 100;

    private final ByteBuffer buffer;
    private final GrowableByteBuffer tableSubscriberEventBuffer;
    private final GrowableByteBuffer globalEventBuffer;
    private final SocketChannel channel;
    private final Object channelWriteLock;
    private final EngineAdapter adapter;
    private final StreamServer server;
    private final ArrayList tableIndexNameMap; // is a list, but used mostly like a map
    private volatile boolean stopped = false;
    private volatile boolean cleanShutdownInProgress = false;
    private volatile boolean receivedHandshake = false;
    private MultiTableWriteHandler multiTableWriteHandler;
    private final StreamSubscriptionsHandler streamSubscriptionsHandler;
    private GlobalSubscriptionsHandler globalSubscriptionsHandler;
    private ThreadedSubscriberEntryWriter subscriptionsWriter;
    private final InetSocketAddress remoteSocketAddress;
    private ClientInfo clientInfo;
    private final AccessControlFilter accessControlFilter;
    private final ReadContinuationCache continuationCache = new ReadContinuationCache();
    private final TimeoutProvider timeoutProvider;

    private class ChunkedResponseHandler implements ChunkedDataHandler {
        private final ByteBuffer buffer;
        private int chunkSizeIndex;

        ChunkedResponseHandler(final ByteBuffer buffer) {
            this.buffer = buffer;
            init();
        }

        private void init() {
            this.chunkSizeIndex = buffer.position();
            buffer.putInt(0);
        }

        @Override
        public void writeChunk() {
            writeChunkImpl();
        }

        private int writeChunkImpl() {
            try {
                final int chunkSize = buffer.position() - chunkSizeIndex - 4;
                buffer.putInt(chunkSizeIndex, chunkSize);
                buffer.flip();
                writeAll(buffer);
                // prepare writing the next chunk
                init();
                return chunkSize;
            } catch (final IOException e) {
                log.error("IOException while writing chunk", e);
                throw new PlovercrestException("IOException while writing chunk", e);
            }
        }

        @Override
        public void writeRemainingData() {
            int chunkSize = writeChunkImpl();
            if (chunkSize > 0) {
                // write end-of-chunks marker (zero size)
                // could be optimized by making it part of the previous write call *if* there is room in the buffer
                chunkSize = writeChunkImpl();
                assert chunkSize == 0;
            }
            buffer.clear();
        }
    }

    public StreamConnection(final StreamServer server,
                            final EngineAdapter adapter,
                            final SocketChannel channel,
                            final InetSocketAddress remoteSocketAddress,
                            final AccessControlFilter accessControlFilter,
                            final TimeoutProvider timeoutFlagProvider) {
        this.accessControlFilter = accessControlFilter;
        this.timeoutProvider = timeoutFlagProvider;
        this.buffer = ByteBuffer.allocate(1024 * 1024);
        this.tableSubscriberEventBuffer = new GrowableByteBuffer(1024);
        this.globalEventBuffer = new GrowableByteBuffer(1024);
        this.server = server;
        this.adapter = adapter;
        this.channel = channel;
        this.remoteSocketAddress = remoteSocketAddress;
        this.channelWriteLock = new Object();
        this.tableIndexNameMap = new ArrayList();
        this.streamSubscriptionsHandler = new StreamSubscriptionsHandler(adapter.getStreamEngine(), this);
        this.clientInfo = new ClientInfo(remoteSocketAddress, "");
    }

    private void readIntoBuffer(final int size) throws IOException, EndOfStreamDuringRead {
        if (buffer.position() != 0) {
            throw new IllegalStateException("Buffer should be in 0 position");
        }
        buffer.limit(size);
        while (buffer.position() < buffer.limit()) {
            if (channel.read(buffer) == -1) {
                throw new EndOfStreamDuringRead();
            }
        }
        buffer.flip();
    }

    private void writeAll(final ByteBuffer buffer) throws IOException {
        // Unfortunately have to lock here because subscriber events happen on a different thread.
        // TODO: switch to async IO and remove locking
        synchronized (channelWriteLock) {
            while (buffer.remaining() > 0) {
                channel.write(buffer);
            }
        }
        buffer.clear();
    }

    private String readString(final int maxLength) throws IOException, EndOfStreamDuringRead {
        readIntoBuffer(4);
        final int length = buffer.getInt();
        if (length < 0 || length > maxLength) {
            throw new PlovercrestException("Cannot handle string of length " + maxLength);
        }
        buffer.clear();
        readIntoBuffer(length);
        // buffer.array() requires a backing array, must be rewritten if we switch to allocateDirect()
        final String value = new String(buffer.array(), 0, length, DefaultPlovercrestRemote.UTF8);
        buffer.clear();
        return value;
    }

    private String readTableName(final int length) throws IOException, EndOfStreamDuringRead {
        if (length < 0 || length > DefaultPlovercrestRemote.MAXIMUM_TABLE_NAME_LENGTH) {
            throw new PlovercrestException("Cannot handle table name with length " + length);
        }
        if (length == 0) {
            return "";
        }
        readIntoBuffer(length);
        // buffer.array() requires a backing array, must be changed if we switch
        // to allocateDirect()
        final String value = new String(buffer.array(), 0, length, DefaultPlovercrestRemote.TABLE_NAME_CHARSET);
        buffer.clear();
        return value;
    }

    @Override
    public void run() {
        try {
            while (!stopped) {
                try {
                    if (buffer.position() != 0) {
                        dump("run() top", buffer);
                        throw new IllegalStateException("Buffer should be empty at this point");
                    }
                    readIntoBuffer(8);
                    // dump("run() after read", buffer);
                    final int action = buffer.getInt();
                    final int tableIndex = buffer.getInt();
                    buffer.clear();
                    final String tableName = determineTableName(tableIndex);

                    // System.err.println("s " + action + " ti " + tableIndex +
                    // " cnt " + count + " sz " +
                    // size);

                    RemoteCommand command;
                    try {
                        command = RemoteCommand.fromTag(action);
                    } catch (final IllegalArgumentException e) {
                        throw new PlovercrestException("Invalid streaming command " + action, e);
                    }
                    if (log.isTraceEnabled()) {
                        log.trace("Processing " + command + " for table " + tableName);
                    }
                    final AccessKind accessKind = AccessControlUtils.getAccessKindForCommand(command);
                    if (accessKind != null && !accessControlFilter.isAccessAllowed(tableName, accessKind, clientInfo)) {
                        writeErrorResponse(new PlovercrestAccessException(accessKind + " access not allowed"));
                        buffer.clear();
                    } else {
                        handleCommand(tableIndex, tableName, command);
                    }
                    // System.err.println("AFTER " + buffer.remaining());
                } catch (final AsynchronousCloseException e) {
                    // async close, this is normal for some client types
                    log.info("Closed stream connection (async close)");
                    stopped = true;
                } catch (final IOException e) {
                    // unclean shutdown, must still stop thread
                    log.warn("Closing stream connection after IOException " + e.getMessage() + " ("
                            + e.getClass().getName() + ")", e);
                    stopped = true;
                } catch (final EndOfStreamDuringRead e) {
                    // clean shutdown
                    log.info("Closed stream connection (end of stream)");
                    stopped = true;
                } catch (final PlovercrestException e) {
                    // unclean shutdown, must still stop thread
                    log.error("Closing stream connection after uncaught plovercrest error: " + e.getMessage(), e);
                    stopped = true;
                }
            }
        } finally {
            close();
        }
    }

    private void logTableEvent(final RemoteCommand command, final String what) {
        if (log.isInfoEnabled()) {
            log.info(command + ": " + what + " [" + clientInfo.getContext() + " " + remoteSocketAddress + "]");
        }
    }

    private void handleCommand(final int tableIndex, final String tableName, final RemoteCommand command)
        throws IOException, EndOfStreamDuringRead {
        // This should be refactored. Maybe register handlers instead of using one big switch statement?
        switch (command) {
            case HANDSHAKE: {
                final String clientContext = readString(DefaultPlovercrestRemote.MAXIMUM_ERROR_MESSAGE_LENGTH);
                buffer.clear();
                handleHandshake(clientContext);
                break;
            }
            case WRITE: {
                readIntoBuffer(8);
                final int count = buffer.getInt();
                final int size = buffer.getInt();
                buffer.clear();
                handleWriteRequest(tableName, count, size);
                break;
            }
            case FLUSH: {
                buffer.clear();
                handleFlush(tableName);
                break;
            }
            case READ_START_OVERLAPPING:
            case READ_START_NON_OVERLAPPING: {
                readIntoBuffer(24);
                final int maxCount = buffer.getInt();
                final long timestamp = buffer.getLong();
                final long endTimestamp = buffer.getLong();
                final ReaderStrategy strategy = ReaderStrategy.fromTag(buffer.getInt());
                buffer.clear();
                handleReadRequest(command, tableName, maxCount, ReaderPosition.createTimestampPosition(timestamp, 0),
                                  endTimestamp, strategy);
                break;
            }
            case READ_ELEMENT_INDEX: {
                readIntoBuffer(24);
                final int maxCount = buffer.getInt();
                final long index = buffer.getLong();
                final long endTimestamp = buffer.getLong();
                final ReaderStrategy strategy = ReaderStrategy.fromTag(buffer.getInt());
                buffer.clear();
                handleReadRequest(command, tableName, maxCount, ReaderPosition.createIndexPosition(index), endTimestamp,
                                  strategy);
                break;
            }
            case READ_CONTINUE_OVERLAPPING:
            case READ_CONTINUE_NON_OVERLAPPING: {
                readIntoBuffer(32);
                final int maxCount = buffer.getInt();
                final int typeTag = buffer.getInt();
                final PositionType positionType = ReaderPositionUtils.typeFromTag(typeTag);
                final long position = buffer.getLong();
                final long endTimestamp = buffer.getLong();
                final int skipCount = buffer.getInt();
                final ReaderStrategy strategy = ReaderStrategy.fromTag(buffer.getInt());
                buffer.clear();
                handleReadRequest(command, tableName, maxCount,
                                  ReaderPosition.create(positionType, position, skipCount), endTimestamp, strategy);
                break;
            }
            case TABLE_INFO: {
                handleTableInfoRequest(tableName);
                break;
            }
            case MULTI_WRITE_INIT: {
                readIntoBuffer(4);
                final int mode = buffer.getInt();
                buffer.clear();
                if (mode < 0 || mode >= WriterMode.values().length) {
                    throw new ProtocolError("illegal writer mode: " + mode);
                }
                handleMultiTableWriteInit(WriterMode.values()[mode]);
                break;
            }
            case MULTI_WRITE: {
                readIntoBuffer(12);
                final int entrySize = buffer.getInt();
                final long timestamp = buffer.getLong();
                buffer.clear();
                handleMultiTableWrite(tableName, entrySize, timestamp);
                break;
            }
            case MULTI_WRITE_ACK: {
                handleMultiTableWriteAck();
                break;
            }
            case CREATE_TABLE_IF_MISSING: {
                readIntoBuffer(4);
                final int size = buffer.getInt();
                if (size < 0 || size > buffer.capacity()) {
                    throw new PlovercrestException("Cannot handle metadata size: " + size);
                }
                buffer.clear();
                readIntoBuffer(size);
                handleCreateTableIfMissing(tableName, buffer);
                buffer.clear();
                break;
            }
            case SUBSCRIBE: {
                readIntoBuffer(16);
                final int tag = buffer.getInt();
                QueryKind queryKind;
                try {
                    queryKind = QueryKind.fromTag(tag);
                } catch (final IllegalArgumentException e) {
                    throw new ProtocolError("Illegal subscription query kind " + tag);
                }
                final long timestampOrIndex = buffer.getLong();
                final int subscriptionId = buffer.getInt();
                buffer.clear();
                handleSubscribe(tableIndex, tableName, new SubscriptionQuery(queryKind, timestampOrIndex),
                                subscriptionId);
                break;
            }
            case UNSUBSCRIBE: {
                readIntoBuffer(4);
                final int subscriptionId = buffer.getInt();
                buffer.clear();
                handleUnsubscribe(subscriptionId);
                break;
            }
            case FREEZE: {
                buffer.clear();
                handleFreeze(tableName);
                break;
            }
            case UNFREEZE: {
                buffer.clear();
                handleUnfreeze(tableName);
                break;
            }
            case SUBSCRIBE_GLOBAL: {
                buffer.clear();
                handleGlobalSubscribe();
                break;
            }
            case LIST: {
                buffer.clear();
                handleList();
                break;
            }
            case CREATE: {
                buffer.clear();
                readIntoBuffer(4);
                final int size = buffer.getInt();
                if (size < 0 || size > buffer.capacity()) {
                    throw new PlovercrestException("Cannot handle metadata size: " + size);
                }
                buffer.clear();
                readIntoBuffer(size);
                handleCreateTable(tableName, buffer);
                buffer.clear();
                break;
            }
            case OPEN: {
                buffer.clear();
                handleOpenTable(tableName);
                break;
            }
            case DELETE: {
                buffer.clear();
                handleDeleteTable(tableName);
                break;
            }
            case RENAME: {
                buffer.clear();
                readIntoBuffer(4);
                final int newTableIndex = buffer.getInt();
                buffer.rewind();
                final String newTableName = determineTableName(newTableIndex);
                buffer.clear();
                handleRenameTable(tableName, newTableName);
                break;
            }
            default:
                throw new RuntimeException("Streaming command not implemented: " + command);
        }
    }

    private void handleHandshake(final String clientContext) throws IOException {
        this.clientInfo = new ClientInfo(remoteSocketAddress, clientContext);
        writeOkResponseWithInt(PROTOCOL_VERSION);
        if (receivedHandshake) {
            log.warn("Received more than one handshake for client: " + clientInfo);
            return;
        }
        receivedHandshake = true;
        if (cleanShutdownInProgress) {
            // if a clean shutdown is in progress already, notify the client now (this should be rare)
            postShutdownEvent();
        }
    }

    private String determineTableName(final int tableIndex) throws IOException, EndOfStreamDuringRead {
        if (tableIndex < 0) {
            throw new ProtocolError("table index negative: " + tableIndex);
        }
        String tableName;
        if (tableIndex < tableIndexNameMap.size()) {
            // known table index
            tableName = tableIndexNameMap.get(tableIndex);
        } else if (tableIndex == tableIndexNameMap.size()) {
            // new table index: read and register new table name
            readIntoBuffer(4);
            final int tableNameLength = buffer.getInt();
            buffer.clear();
            tableName = readTableName(tableNameLength);
            tableIndexNameMap.add(tableName);
        } else {
            throw new ProtocolError("table index out of bounds: " + tableIndex);
        }
        return tableName;
    }

    private void close() {
        try {
            if (channel.isOpen()) {
                channel.close();
            }
        } catch (final IOException e) {
            log.error("Error closing open channel", e);
        }
        server.deregister(this);
        if (multiTableWriteHandler != null) {
            multiTableWriteHandler.close();
            multiTableWriteHandler = null;
        }
        streamSubscriptionsHandler.close();
        if (globalSubscriptionsHandler != null) {
            globalSubscriptionsHandler.close();
            globalSubscriptionsHandler = null;
        }
        if (subscriptionsWriter != null) {
            subscriptionsWriter.close();
            subscriptionsWriter = null;
        }
    }

    private void handleSubscribe(final int tableIndex, final String tableName, final SubscriptionQuery query,
                                 final int subscriptionId) throws IOException {
        streamSubscriptionsHandler.add(tableName, query, subscriptionId);
    }

    private void handleUnsubscribe(final int subscriptionId) throws IOException {
        streamSubscriptionsHandler.remove(subscriptionId);
    }

    private void handleGlobalSubscribe() {
        if (globalSubscriptionsHandler == null) {
            globalSubscriptionsHandler = new GlobalSubscriptionsHandler(this, adapter.getStreamEngine());
        }
    }

    private void dump(final String pre, final ByteBuffer buf) {
        System.err.println(pre + " { rem " + buf.remaining() + " / pos " + buf.position() + " / lim " + buf.limit()
                + "}");
    }

    private void prepareResponse(final RemoteResponse responseType) {
        if (buffer.position() != 0) {
            throw new IllegalStateException("Buffer should be in 0 position");
        }
        buffer.putInt(responseType.getTag());
    }

    private void writeErrorResponse(final PlovercrestException e) throws IOException {
        writeMultiErrorResponse(Collections.singleton(new IndexedError(-1, e)));
    }

    private void writeMultiErrorResponse(final Collection errors) throws IOException {
        // handle errors by serialize the exceptions in order to recreate on the client side
        assert !errors.isEmpty();
        buffer.clear();
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_ERROR);
        buffer.putInt(errors.size());
        for (final IndexedError err : errors) {
            buffer.putInt(err.getIndex());
            final PlovercrestException e = err.getException();
            buffer.putInt(e.errorType().getTag());
            final byte[] data = e.errorData().getBytes(DefaultPlovercrestRemote.ERROR_MESSAGE_CHARSET);
            buffer.putInt(data.length);
            buffer.put(data);
        }
        buffer.flip();
        writeAll(buffer);
    }

    private void handleReadRequest(final RemoteCommand readCommand, final String tableName, final int count,
                                   final ReaderPosition position, final long endTimestampOrIndex,
                                   final ReaderStrategy strategy) throws IOException {
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        final ChunkedResponseHandler dataHandler = new ChunkedResponseHandler(buffer);
        final AtomicBoolean timedOut;
        switch (strategy) {
            case GREEDY:
                // Never time out greedy read requests
                timedOut = new AtomicBoolean();
                break;
            case NON_GREEDY:
                timedOut = timeoutProvider.createTimeoutFlag(READ_TIMEOUT_MILLISECONDS, TimeUnit.MILLISECONDS);
                break;
            default:
                throw new IllegalStateException("Unknown strategy " + strategy);
        }
        final ReadRequest readRequest =
                new ReadRequest(continuationCache, tableName, count, buffer, dataHandler, timedOut);
        switch (readCommand) {
            case READ_START_NON_OVERLAPPING:
                adapter.read(readRequest, RangeQuery.createNonOverlappingTimeRange(position.getTimestampOrIndex(),
                                                                                   endTimestampOrIndex));
                break;
            case READ_START_OVERLAPPING:
                adapter.read(readRequest, RangeQuery.createOverlappingTimeRange(position.getTimestampOrIndex(),
                                                                                endTimestampOrIndex));
                break;
            case READ_CONTINUE_NON_OVERLAPPING:
                adapter.readContinue(readRequest, RangeType.NON_OVERLAPPING_TIME_RANGE, endTimestampOrIndex, position);
                break;
            case READ_CONTINUE_OVERLAPPING:
                adapter.readContinue(readRequest, RangeType.OVERLAPPING_TIME_RANGE, endTimestampOrIndex, position);
                break;
            case READ_ELEMENT_INDEX:
                adapter.read(readRequest,
                             RangeQuery.createIndexRange(position.getTimestampOrIndex(), endTimestampOrIndex));
                break;
            default:
                throw new IllegalStateException("Not a known read command: " + readCommand);
        }
        // data chunks are already written, but still have to write the trailer
        buffer.flip();
        writeAll(buffer);
    }

    private void handleWriteRequest(final String tableName, final int count, final int size)
        throws IOException, EndOfStreamDuringRead {
        if (buffer.position() != 0) {
            throw new IllegalStateException("Buffer should be empty");
        }
        final SizeLimitedByteBufferReader reader = new SizeLimitedByteBufferReader(buffer, size, channel);

        // System.err.println("c +" + buffer.remaining());
        final ArrayList errors = new ArrayList<>();
        try {
            for (int i = 0; i < count; i++) {
                final int entrySize = reader.readInt();
                final long timestamp = reader.readLong();
                try {
                    adapter.write(tableName, timestamp, reader.readSlice(entrySize));
                } catch (final PlovercrestException e) {
                    log.error(e.getMessage(), e);
                    errors.add(new IndexedError(i, e));
                }
            }
            if (reader.remaining() != 0) {
                throw new IllegalStateException("Should have read all data by now");
            }
            buffer.clear();
            // System.err.println();
            // dump("after handleRequest() ", buffer);

            // write ack value for client
            if (errors.isEmpty()) {
                writeSimpleAckResponse();
            } else {
                writeMultiErrorResponse(errors);
            }
        } catch (final PlovercrestException e) {
            writeErrorResponse(e);
            buffer.clear();
        }
    }

    private void handleTableInfoRequest(final String tableName) throws IOException {
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        final int limit = buffer.position() + 8 * 5 + 1;
        buffer.limit(limit);
        adapter.readStreamInfo(tableName, buffer);
        if (buffer.remaining() > 0) {
            throw new IllegalStateException("Not enough table info data read");
        }
        buffer.flip();
        writeAll(buffer);
    }

    private void handleMultiTableWriteInit(final WriterMode mode) throws IOException, EndOfStreamDuringRead {
        // log.debug("handleMultiTableWriteInit size=" + objectSize);
        if (multiTableWriteHandler != null) {
            multiTableWriteHandler.close();
        }
        multiTableWriteHandler = new MultiTableWriteHandler(adapter, mode);
        // TODO: catch errors and send back ok/error response
        buffer.clear();
    }

    private void handleMultiTableWrite(final String tableName, final int entrySize, final long timestamp)
        throws IOException, EndOfStreamDuringRead {
        readIntoBuffer(entrySize);
        if (multiTableWriteHandler == null) {
            throw new ProtocolError("multi-table write not initialized");
        } else {
            multiTableWriteHandler.write(tableName, timestamp, buffer);
        }
        buffer.clear();
    }

    private void handleMultiTableWriteAck() throws IOException {
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        buffer.putInt(ACK);
        buffer.flip();
        writeAll(buffer);
    }

    private void handleCreateTableIfMissing(final String tableName, final ByteBuffer data) {
        if (adapter.open(tableName) == null) {
            logTableEvent(RemoteCommand.CREATE_TABLE_IF_MISSING, tableName);
            final int fanoutDistance = data.getInt();
            final TableMetadata metadata = TableMetadataUtils.read(data);
            adapter.create(new TableInfo(tableName, fanoutDistance, metadata));
        }
    }

    private void handleCreateTable(final String tableName, final ByteBuffer data) throws IOException {
        logTableEvent(RemoteCommand.CREATE, tableName);
        final int fanoutDistance = data.getInt();
        final TableMetadata metadata = TableMetadataUtils.read(data);
        try {
            adapter.create(new TableInfo(tableName, fanoutDistance, metadata));
        } catch (final PlovercrestException e) {
            buffer.clear();
            writeErrorResponse(e);
            return;
        }
        buffer.clear();
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        buffer.flip();
        writeAll(buffer);
    }

    private void handleOpenTable(final String tableName) throws IOException {
        try {
            final TableMetadata metadata = adapter.getMetadata(tableName);
            prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
            if (metadata == null) {
                buffer.putInt(0);
            } else {
                buffer.putInt(PlovercrestConstants.DEFAULT_FANOUT_DISTANCE); // backwards compatibility
                putMetaDataWithSize(buffer, metadata);
            }
            buffer.flip();
            writeAll(buffer);
        } catch (final PlovercrestException e) {
            writeErrorResponse(e);
        }
    }

    private void handleDeleteTable(final String tableName) throws IOException {
        logTableEvent(RemoteCommand.DELETE, tableName);
        final boolean result = adapter.delete(tableName);
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        final byte booleanResponse = (byte) (result ? 1 : 0);
        buffer.put(booleanResponse);
        buffer.flip();
        writeAll(buffer);
    }

    private void handleFlush(final String tableName) throws IOException {
        adapter.flush(tableName);
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        buffer.flip();
        writeAll(buffer);
    }

    private void handleRenameTable(final String oldName, final String newName) throws IOException {
        logTableEvent(RemoteCommand.RENAME, oldName + " -> " + newName);
        final boolean renamed = adapter.rename(oldName, newName);
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        final byte booleanResponse = (byte) (renamed ? 1 : 0);
        buffer.put(booleanResponse);
        buffer.flip();
        writeAll(buffer);
    }

    private void handleFreeze(final String tableName) throws IOException {
        logTableEvent(RemoteCommand.FREEZE, tableName);
        adapter.freeze(tableName);
        writeSimpleAckResponse();
    }

    private void writeOkResponseWithInt(final int ack) throws IOException {
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        buffer.putInt(ack);
        buffer.flip();
        writeAll(buffer);
    }

    private void writeSimpleAckResponse() throws IOException {
        writeOkResponseWithInt(ACK);

    }

    private void handleUnfreeze(final String tableName) throws IOException {
        logTableEvent(RemoteCommand.UNFREEZE, tableName);
        adapter.unfreeze(tableName);
        writeSimpleAckResponse();
    }

    private static void putMetaDataWithSize(final ByteBuffer out, final TableMetadata metaData) {
        final int startPos = out.position();
        out.putInt(0);
        TableMetadataUtils.write(metaData, out);
        final int size = out.position() - startPos - 4;
        out.putInt(startPos, size);
    }

    private void handleList() throws IOException {
        prepareResponse(RemoteResponse.COMMAND_RESPONSE_OK);
        final Collection infos = adapter.list();
        buffer.putInt(infos.size());
        for (final StreamInfo info : infos) {
            final byte[] data = info.getName().getBytes(DefaultPlovercrestRemote.TABLE_NAME_CHARSET);
            if (data.length > DefaultPlovercrestRemote.MAXIMUM_TABLE_NAME_LENGTH) {
                throw new ProtocolError("Table name too long (" + data.length + ")");
            }
            if (buffer.remaining() < data.length + 5) {
                // no room for next element, must flush
                buffer.flip();
                writeAll(buffer);
            }
            buffer.putInt(data.length);
            buffer.put(data);
            buffer.put((byte) (info.isMaterialized() ? 1 : 0));

        }
        buffer.flip();
        writeAll(buffer);
    }

    public void doStop() {
        if (stopped) {
            return;
        }
        try {
            channel.close();
        } catch (final IOException e) {
            throw new PlovercrestException("Failed to close stream", e);
        }
        stopped = true;
    }

    private void flipAndWriteSubscriberEvent(final ByteBuffer tmpBuffer) {
        // This is called on a different thread from other write calls.
        // It should be thread safe since we use a separate ByteBuffer, and every response
        // message is written through a synchronized writeAll() call. Still, an async socket
        // interface would be better.
        if (stopped) {
            return;
        }
        tmpBuffer.flip();
        try {
            writeAll(tmpBuffer);
        } catch (final IOException e) {
            log.error("Error while posting subscriber event: " + e.getMessage(), e);
            stopped = true;
            // closing the channel should cause the main loop to terminate quickly
            if (channel.isOpen()) {
                try {
                    channel.close();
                } catch (final IOException e1) {
                    log.error("Unexpected error closing channel", e1);
                }
            }
        }
    }

    private void writeEntries(final SubscriberEntryCollection c) {
        if (stopped) {
            return;
        }
        try {
            for (final List entries : c.getEntriesByTableIndex().values()) {
                assert entries.size() > 0;
                if (entries.size() == 1) {
                    final SubscriberEntry entry = entries.get(0);
                    final ByteBuffer tmpBuffer = ByteBuffer.allocate(entry.data.length + 28);
                    tmpBuffer.putInt(RemoteResponse.SUBSCRIBE_ENTRY.getTag());
                    tmpBuffer.putInt(entry.subscriptionId);
                    tmpBuffer.putLong(entry.entryIndex);
                    tmpBuffer.putLong(entry.timestamp);
                    tmpBuffer.putInt(entry.data.length);
                    tmpBuffer.put(entry.data);
                    tmpBuffer.flip();
                    writeAll(tmpBuffer);
                } else {
                    int totalEntriesSize = 0;
                    for (final SubscriberEntry entry : entries) {
                        totalEntriesSize += 8 + entry.data.length;
                    }
                    final int totalBufferSize = 24 + totalEntriesSize;
                    final ByteBuffer tmpBuffer = ByteBuffer.allocate(totalBufferSize);
                    final SubscriberEntry firstEntry = entries.get(0);
                    tmpBuffer.putInt(RemoteResponse.SUBSCRIBE_ENTRY_CHUNK.getTag());
                    tmpBuffer.putInt(firstEntry.subscriptionId);
                    tmpBuffer.putLong(firstEntry.entryIndex);
                    tmpBuffer.putInt(entries.size());
                    tmpBuffer.putInt(totalEntriesSize);
                    for (int i = 0; i < entries.size(); i++) {
                        final SubscriberEntry entry = entries.get(i);
                        assert entry.subscriptionId == firstEntry.subscriptionId;
                        assert entry.entryIndex == firstEntry.entryIndex + i;
                        tmpBuffer.putLong(entry.timestamp);
                        tmpBuffer.put(entry.data);
                    }
                    tmpBuffer.flip();
                    writeAll(tmpBuffer);
                }
            }
        } catch (final IOException e) {
            log.error("Error while posting subscriber events: " + e.getMessage(), e);
            stopped = true;
            // closing the channel should cause the main loop to terminate quickly
            if (channel.isOpen()) {
                try {
                    channel.close();
                } catch (final IOException e1) {
                    log.error("Unexpected error closing channel", e1);
                }
            }
        }
    }

    void postEntryToSubscriber(final int subscriptionId, final long entryIndex, final long timestamp,
                               final ByteBuffer bytes) {
        ensureSubscriberWriterIsCreated();
        final byte[] data = new byte[bytes.remaining()];
        bytes.get(data);
        subscriptionsWriter.put(new SubscriberEntry(subscriptionId, entryIndex, timestamp, data));
    }

    private void ensureSubscriberWriterIsCreated() {
        if (subscriptionsWriter == null) {
            subscriptionsWriter = new ThreadedSubscriberEntryWriter(new SubscriberWriteHandler() {
                @Override
                public void writeEntryCollection(final SubscriberEntryCollection c) {
                    writeEntries(c);
                }

                @Override
                public void writeStandaloneEvent(final StandaloneSubscriberEvent event) {
                    flipAndWriteSubscriberEvent(event.serialize());
                }
            });
        }
    }

    void postLastValidTimestampToSubscriber(final int subscriptionId, final long entryIndex,
                                            final long lastValidTimestamp) {
        ensureSubscriberWriterIsCreated();
        subscriptionsWriter.put(new SubscriberLastValidTimestampEvent(subscriptionId, entryIndex, lastValidTimestamp));
    }

    void postAckToSubscriber(final int subscriptionId, final long entryIndex) {
        ensureSubscriberWriterIsCreated();
        subscriptionsWriter.put(new SubscriberAck(subscriptionId, entryIndex));
    }

    void postTableCreatedToSubscriber(final int subscriptionId) {
        final ByteBuffer tmpBuffer = tableSubscriberEventBuffer.getBufferWithLimit(8);
        tmpBuffer.putInt(RemoteResponse.SUBSCRIBE_TABLE_CREATED.getTag());
        tmpBuffer.putInt(subscriptionId);
        flipAndWriteSubscriberEvent(tmpBuffer);
    }

    void postTableFrozenToSubscriber(final int subscriptionId) {
        ensureSubscriberWriterIsCreated();
        subscriptionsWriter.put(new SubscriberFrozenEvent(subscriptionId));
    }

    private static void putTableName(final ByteBuffer buf, final byte[] data) {
        buf.putInt(data.length);
        buf.put(data);
    }

    // We use the full table names for global events. Cannot use table indexes because they are decided
    // by the client.  Should not be a problem because global table events are not very frequent.
    // Note that global table events are received on a different thread from table subscriber events,
    // so a different ByteBuffer has to be used (until we switch to async socket IO).

    private void postGlobalTableNameEvent(final RemoteResponse response, final Collection names) {
        int totalLength = 8; // response + number of names
        final ArrayList nameData = new ArrayList();
        for (final String name : names) {
            final byte[] data = name.getBytes(DefaultPlovercrestRemote.TABLE_NAME_CHARSET);
            nameData.add(data);
            totalLength += 4 + data.length;
        }
        final ByteBuffer tmpBuffer = globalEventBuffer.getBufferWithLimit(totalLength);
        tmpBuffer.putInt(response.getTag());
        tmpBuffer.putInt(nameData.size());
        for (final byte[] data : nameData) {
            putTableName(tmpBuffer, data);
        }
        flipAndWriteSubscriberEvent(tmpBuffer);
    }

    void postGlobalTableCreated(final String tableName) {
        postGlobalTableNameEvent(RemoteResponse.GLOBAL_TABLE_CREATED, Collections.singleton(tableName));
    }

    void postGlobalTableDeleted(final String tableName) {
        postGlobalTableNameEvent(RemoteResponse.GLOBAL_TABLE_DELETED, Collections.singleton(tableName));
    }

    void postGlobalResetTableNames(final Collection tableNames) {
        postGlobalTableNameEvent(RemoteResponse.GLOBAL_RESET_TABLES, tableNames);
    }

    private void postShutdownEvent() {
        final ByteBuffer tmpBuffer = ByteBuffer.allocate(4);
        tmpBuffer.putInt(RemoteResponse.GLOBAL_SHUTDOWN_EVENT.getTag());
        flipAndWriteSubscriberEvent(tmpBuffer);
    }

    void startCleanShutdown() {
        if (cleanShutdownInProgress) {
            return;
        }
        cleanShutdownInProgress = true;
        if (!receivedHandshake) {
            // Sending events before the initial heartbeat is not handled by clients, so delay it.
            return;
        }
        postShutdownEvent();
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy