content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *   http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   * KIND, either express or implied.  See the License for the
   * specific language governing permissions and limitations
   * under the License.
   */
  package org.apache.asterix.replication.management;
  
  import java.io.BufferedReader;
  import java.io.IOException;
  import java.io.InputStreamReader;
  import java.nio.ByteBuffer;
  import java.nio.channels.AsynchronousCloseException;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.Executor;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.LinkedBlockingQueue;
  
  import org.apache.asterix.common.api.INcApplicationContext;
  import org.apache.asterix.common.config.ReplicationProperties;
  import org.apache.asterix.common.exceptions.ReplicationException;
  import org.apache.asterix.common.replication.IReplicationDestination;
  import org.apache.asterix.common.replication.IReplicationManager;
  import org.apache.asterix.common.transactions.ILogRecord;
  import org.apache.asterix.common.transactions.LogType;
  import org.apache.asterix.replication.api.ReplicationDestination;
  import org.apache.asterix.replication.logging.ReplicationLogBuffer;
  import org.apache.asterix.replication.logging.TxnAckTracker;
  import org.apache.asterix.replication.logging.TxnLogReplicator;
  import org.apache.asterix.replication.messaging.ReplicateLogsTask;
  import org.apache.asterix.replication.messaging.ReplicationProtocol;
  import org.apache.hyracks.api.network.ISocketChannel;
  import org.apache.logging.log4j.LogManager;
  import org.apache.logging.log4j.Logger;
  
  public class LogReplicationManager {
  
      private static final Logger LOGGER = LogManager.getLogger();
      private final LinkedBlockingQueue<ReplicationLogBuffer> emptyLogBuffersQ;
      private final LinkedBlockingQueue<ReplicationLogBuffer> pendingFlushLogBuffersQ;
      private final ByteBuffer txnLogsBatchSizeBuffer = ByteBuffer.allocate(Integer.BYTES);
      private final Map<ReplicationDestination, ISocketChannel> destinations = new HashMap<>();
      private final IReplicationManager replicationManager;
      private final Executor executor;
      private final TxnAckTracker ackTracker = new TxnAckTracker();
      private final Set<ISocketChannel> failedSockets = new HashSet<>();
      private final Object transferLock = new Object();
      private final INcApplicationContext appCtx;
      private final int logPageSize;
      private final int logBatchSize;
      private ReplicationLogBuffer currentTxnLogBuffer;
      private ISocketChannel[] destSockets;
  
      public LogReplicationManager(INcApplicationContext appCtx, IReplicationManager replicationManager) {
          this.appCtx = appCtx;
          this.replicationManager = replicationManager;
          final ReplicationProperties replicationProperties = appCtx.getReplicationProperties();
          logPageSize = replicationProperties.getLogBufferPageSize();
          logBatchSize = replicationProperties.getLogBatchSize();
          executor = appCtx.getThreadExecutor();
          emptyLogBuffersQ = new LinkedBlockingQueue<>();
          pendingFlushLogBuffersQ = new LinkedBlockingQueue<>();
          initBuffers(replicationProperties.getLogBufferNumOfPages());
          TxnLogReplicator txnlogReplicator = new TxnLogReplicator(emptyLogBuffersQ, pendingFlushLogBuffersQ);
          ((ExecutorService) executor).submit(txnlogReplicator);
      }
  
      private void initBuffers(int buffers) {
          for (int i = 0; i < buffers; i++) {
              emptyLogBuffersQ.add(new ReplicationLogBuffer(this, logPageSize, logBatchSize));
          }
          try {
              getAndInitNewPage();
          } catch (InterruptedException e) {
              Thread.currentThread().interrupt();
              throw new ReplicationException(e);
          }
      }
  
      public void register(ReplicationDestination dest) {
          synchronized (transferLock) {
              synchronized (destinations) {
                  if (destinations.containsKey(dest)) {
                      return;
                  }
                  LOGGER.info(() -> "register " + dest);
                  ISocketChannel socketChannel = dest.getLogReplicationChannel(appCtx);
                  handshake(dest, socketChannel);
                  destinations.put(dest, socketChannel);
                  failedSockets.remove(socketChannel);
                  destSockets = destinations.values().toArray(new ISocketChannel[0]);
              }
          }
      }
  
      public void unregister(IReplicationDestination dest) {
          synchronized (transferLock) {
              synchronized (destinations) {
                  if (!destinations.containsKey(dest)) {
                      return;
                  }
                  LOGGER.info(() -> "unregister " + dest);
                  ackTracker.unregister(dest);
                  ISocketChannel destSocket = destinations.remove(dest);
                  failedSockets.remove(destSocket);
                  destSockets = destinations.values().toArray(new ISocketChannel[0]);
                  endReplication(destSocket);
              }
          }
      }
  
      public void replicate(ILogRecord logRecord) throws InterruptedException {
          if (logRecord.getLogType() == LogType.JOB_COMMIT || logRecord.getLogType() == LogType.ABORT) {
              synchronized (destinations) {
                  ackTracker.track(logRecord, new HashSet<>(destinations.keySet()));
              }
          }
          appendToLogBuffer(logRecord);
      }
  
      public void transferBatch(final ByteBuffer buffer) {
          // prepare the batch size buffer
          txnLogsBatchSizeBuffer.clear();
          txnLogsBatchSizeBuffer.putInt(buffer.remaining());
          txnLogsBatchSizeBuffer.flip();
  
          buffer.mark();
          synchronized (transferLock) {
              if (destSockets != null) {
                  for (ISocketChannel replicaSocket : destSockets) {
                      try {
                          // send batch size then the batch itself
                          NetworkingUtil.transferBufferToChannel(replicaSocket, txnLogsBatchSizeBuffer);
                          NetworkingUtil.transferBufferToChannel(replicaSocket, buffer);
                      } catch (IOException e) {
                          handleFailure(replicaSocket, e);
                      } finally {
                          txnLogsBatchSizeBuffer.position(0);
                          buffer.reset();
                      }
                  }
              }
          }
          // move the buffer position to the sent limit
          buffer.position(buffer.limit());
      }
  
      public int getLogPageSize() {
          return logPageSize;
      }
  
      private synchronized void appendToLogBuffer(ILogRecord logRecord) throws InterruptedException {
          if (!currentTxnLogBuffer.hasSpace(logRecord)) {
              currentTxnLogBuffer.setFull(true);
              if (logRecord.getLogSize() > logPageSize) {
                  getAndInitNewLargePage(logRecord.getLogSize());
              } else {
                  getAndInitNewPage();
              }
          }
          currentTxnLogBuffer.append(logRecord);
      }
  
      private void getAndInitNewPage() throws InterruptedException {
          currentTxnLogBuffer = null;
          while (currentTxnLogBuffer == null) {
              currentTxnLogBuffer = emptyLogBuffersQ.take();
          }
          currentTxnLogBuffer.reset();
          pendingFlushLogBuffersQ.add(currentTxnLogBuffer);
      }
  
      private void getAndInitNewLargePage(int pageSize) {
          // for now, alloc a new buffer for each large page
          currentTxnLogBuffer = new ReplicationLogBuffer(this, pageSize, logBatchSize);
          pendingFlushLogBuffersQ.add(currentTxnLogBuffer);
      }
  
      private void handshake(ReplicationDestination dest, ISocketChannel socketChannel) {
          final String nodeId = appCtx.getServiceContext().getNodeId();
          final ReplicateLogsTask task = new ReplicateLogsTask(nodeId);
          ReplicationProtocol.sendTo(socketChannel, task, null);
          executor.execute(new TxnAckListener(dest, socketChannel));
      }
  
      private void endReplication(ISocketChannel socketChannel) {
          if (socketChannel.getSocketChannel().isConnected()) {
              // end log replication (by sending a dummy log with a single byte)
              final ByteBuffer endLogRepBuffer = ReplicationProtocol.getEndLogReplicationBuffer();
              try {
                  NetworkingUtil.transferBufferToChannel(socketChannel, endLogRepBuffer);
              } catch (IOException e) {
                  LOGGER.warn("Failed to end txn log", e);
              }
          }
      }
  
      private synchronized void handleFailure(ISocketChannel replicaSocket, Exception e) {
          if (failedSockets.contains(replicaSocket)) {
              return;
          }
          LOGGER.debug("Replica failed", e);
          failedSockets.add(replicaSocket);
          Optional<ReplicationDestination> socketDest = destinations.entrySet().stream()
                  .filter(entry -> entry.getValue().equals(replicaSocket)).map(Map.Entry::getKey).findFirst();
          if (socketDest.isPresent()) {
              ReplicationDestination dest = socketDest.get();
              LOGGER.error("replica at {} failed", dest);
              replicationManager.notifyFailure(dest, e);
          }
      }
  
      private class TxnAckListener implements Runnable {
          private final ReplicationDestination dest;
          private final ISocketChannel replicaSocket;
  
          TxnAckListener(ReplicationDestination dest, ISocketChannel replicaSocket) {
              this.dest = dest;
              this.replicaSocket = replicaSocket;
          }
  
          @Override
          public void run() {
              Thread.currentThread().setName("TxnAckListener (" + dest + ")");
              LOGGER.info("Started listening on socket: {}", dest);
              try (BufferedReader incomingResponse = new BufferedReader(
                      new InputStreamReader(replicaSocket.getSocketChannel().socket().getInputStream()))) {
                  while (true) {
                      final String response = incomingResponse.readLine();
                      if (response == null) {
                          handleFailure(replicaSocket, new IOException("Unexpected response from replica " + dest));
                          break;
                      }
                      // read ACK
                      final int txnId = ReplicationProtocol.getTxnIdFromLogAckMessage(response);
                      ackTracker.ack(txnId, dest);
                  }
              } catch (AsynchronousCloseException e) {
                  LOGGER.debug(() -> "Stopped listening on socket:" + dest, e);
              } catch (Exception e) {
                  handleFailure(replicaSocket, e);
              }
          }
      }
  }
  