connectiontablenio.java

来自「JGRoups源码」· Java 代码 · 共 1,494 行 · 第 1/4 页

               {
                  SelectionKey key = (SelectionKey) i.next();
                  i.remove();
                  // Do partial read and handle call back
                  Connection conn = (Connection) key.attachment();
                  try
                  {
                     if (conn.getSocketChannel().isOpen())
                        readOnce(conn);
                     else
                     {  // socket connection is already closed, clean up connection state
                        conn.closed();
                     }
                  } catch (IOException e)
                  {
                     if (LOG.isTraceEnabled()) LOG.trace("Read operation on socket failed" , e);
                     // The connection must be bad, cancel the key, close socket, then
                     // remove it from table!
                     key.cancel();
                     conn.destroy();
                     conn.closed();
                  }
               }
            }

            // Now we look at the connection queue to get any new connections added
            Object o;
            try
            {
               o = QUEUE.poll(0); // get a connection
            } catch (InterruptedException e)
            {
               if (LOG.isInfoEnabled()) LOG.info("Thread ("+Thread.currentThread().getName() +") was interrupted while polling queue" ,e);
               // We must give up
               continue;
            }
            if (null == o)
               continue;
            if (o instanceof Shutdown) {     // shutdown command?
               try {
                  SELECTOR.close();
               } catch(IOException e) {
                  if (LOG.isInfoEnabled()) LOG.info("Read selector close operation failed" , e);
               }
               return;                       // stop reading
            }
            Connection conn = (Connection) o;// must be a new connection
            SocketChannel sc = conn.getSocketChannel();
            try
            {
               sc.register(SELECTOR, SelectionKey.OP_READ, conn);
            } catch (ClosedChannelException e)
            {
               if (LOG.isInfoEnabled()) LOG.info("Socket channel was closed while we were trying to register it to selector" , e);
               // Channel becomes bad. The connection must be bad,
               // close socket, then remove it from table!
               conn.destroy();
               conn.closed();
            }
         }   // end of the while true loop
      }

      private void readOnce(Connection conn)
         throws IOException
      {
         ConnectionReadState readState = conn.getReadState();
         if (!readState.isHeadFinished())
         {  // a brand new message coming or header is not completed
            // Begin or continue to read header
            int size = readHeader(conn);
            if (0 == size)
            {  // header is not completed
               return;
            }
         }
         // Begin or continue to read body
         if (readBody(conn) > 0)
         { // not finish yet
            return;
         }
         Address addr = conn.getPeerAddress();
         ByteBuffer buf = readState.getReadBodyBuffer();
         // Clear status
         readState.bodyFinished();
         // Assign worker thread to execute call back
         try
         {
            connectTable.runRequest(addr, buf);
         } catch (InterruptedException e)
         {
            // Cannot do call back, what can we do?
            // Give up handling the message then
            LOG.error("Thread ("+Thread.currentThread().getName() +") was interrupted while assigning executor to process read request" , e);
         }
      }

      /**
       * Read message header from channel. It doesn't try to complete. If there is nothing in
       * the channel, the method returns immediately.
       *
       * @param conn The connection
       * @return 0 if header hasn't been read completely, otherwise the size of message body
       * @throws IOException
       */
      private int readHeader(Connection conn)
         throws IOException
      {
         ConnectionReadState readState = conn.getReadState();
         ByteBuffer headBuf = readState.getReadHeadBuffer();

         SocketChannel sc = conn.getSocketChannel();
         while (headBuf.remaining() > 0)
         {
            int num = sc.read(headBuf);
            if (-1 == num)
            {// EOS
               throw new IOException("Peer closed socket");
            }
            if (0 == num) // no more data
               return 0;
         }
         // OK, now we get the whole header, change the status and return message size
         return readState.headFinished();
      }

      /**
       * Read message body from channel. It doesn't try to complete. If there is nothing in
       * the channel, the method returns immediately.
       *
       * @param conn The connection
       * @return remaining bytes for the message
       * @throws IOException
       */
      private int readBody(Connection conn)
         throws IOException
      {
         ByteBuffer bodyBuf = conn.getReadState().getReadBodyBuffer();

         SocketChannel sc = conn.getSocketChannel();
         while (bodyBuf.remaining() > 0)
         {
            int num = sc.read(bodyBuf);
            if (-1 == num) // EOS
               throw new IOException("Couldn't read from socket as peer closed the socket");
            if (0 == num) // no more data
               return bodyBuf.remaining();
         }
         // OK, we finished reading the whole message! Flip it (not necessary though)
         bodyBuf.flip();
         return 0;
      }
   }

   private class ExecuteTask implements Runnable {
      Address m_addr = null;
      ByteBuffer m_buf = null;

      public ExecuteTask(Address addr, ByteBuffer buf)
      {
         m_addr = addr;
         m_buf = buf;
      }

      public void run()
      {
         receive(m_addr, m_buf.array(), m_buf.arrayOffset(), m_buf.limit());
      }
   }

   private class ConnectionReadState {
      private final Connection m_conn;

      // Status for receiving message
      private boolean m_headFinished = false;
      private ByteBuffer m_readBodyBuf = null;
      private final ByteBuffer m_readHeadBuf = ByteBuffer.allocate(Connection.HEADER_SIZE);

      public ConnectionReadState(Connection conn)
      {
         m_conn = conn;
      }

      ByteBuffer getReadBodyBuffer()
      {
         return m_readBodyBuf;
      }

      ByteBuffer getReadHeadBuffer()
      {
         return m_readHeadBuf;
      }

      void bodyFinished()
      {
         m_headFinished = false;
         m_readHeadBuf.clear();
         m_readBodyBuf = null;
         m_conn.updateLastAccessed();
      }

      /**
       * Status change for finishing reading the message header (data already in buffer)
       *
       * @return message size
       */
      int headFinished()
      {
         m_headFinished = true;
         m_readHeadBuf.flip();
         int messageSize = m_readHeadBuf.getInt();
         m_readBodyBuf = ByteBuffer.allocate(messageSize);
         m_conn.updateLastAccessed();
         return messageSize;
      }

      boolean isHeadFinished()
      {
         return m_headFinished;
      }
   }

   class Connection extends ConnectionTable.Connection {
      private SocketChannel sock_ch = null;
      private WriteHandler m_writeHandler;
      private SelectorWriteHandler m_selectorWriteHandler;
      private final ConnectionReadState m_readState;

      private static final int HEADER_SIZE = 4;
      final ByteBuffer headerBuffer = ByteBuffer.allocate(HEADER_SIZE);

      Connection(SocketChannel s, Address peer_addr)
      {
         super(s.socket(), peer_addr);
         sock_ch = s;
         m_readState = new ConnectionReadState(this);
          is_running=true;
      }

      private ConnectionReadState getReadState()
      {
         return m_readState;
      }

      private void setupWriteHandler(WriteHandler hdlr)
      {
         m_writeHandler = hdlr;
         m_selectorWriteHandler = hdlr.add(sock_ch);
      }

//      void destroy()
//      {
//         closeSocket();
//      }

      void doSend(byte[] buffie, int offset, int length) throws Exception
      {
         FutureResult result = new FutureResult();
         m_writeHandler.write(sock_ch, ByteBuffer.wrap(buffie, offset, length), result, m_selectorWriteHandler);
         Exception ex = result.getException();
         if (ex != null)
         {
            if (LOG.isErrorEnabled())
               LOG.error("failed sending message", ex);
            if (ex.getCause() instanceof IOException)
               throw (IOException) ex.getCause();
            throw ex;
         }
         result.get();
      }


      SocketChannel getSocketChannel()
      {
         return sock_ch;
      }

      void closeSocket()
      {

         if (sock_ch != null)
         {
            try
            {
               if(sock_ch.isConnected() && sock_ch.isOpen()) {
                  sock_ch.close();
               }
            }
            catch (Exception e)
            {
               log.error("error closing socket connection", e);
            }
            sock_ch = null;
         }
      }


      void closed()
      {
         Address peerAddr = getPeerAddress();
         synchronized (conns)
         {
            conns.remove(peerAddr);
         }
         notifyConnectionClosed(peerAddr);
      }
   }


   /**
    * Handle writing to non-blocking NIO connection.
    */
   private static class WriteHandler implements Runnable {
      // Create a queue for write requests
      private final LinkedQueue QUEUE = new LinkedQueue();

      private final Selector SELECTOR = initSelector();
      private int m_pendingChannels;                 // count of the number of channels that have pending writes
      // note that this variable is only accessed by one thread.

      // allocate and reuse the header for all buffer write operations
      private ByteBuffer m_headerBuffer = ByteBuffer.allocate(Connection.HEADER_SIZE);


      Selector initSelector() {
         try
         {
            return SelectorProvider.provider().openSelector();
         }
         catch (IOException e)
         {
            if (LOG.isErrorEnabled()) LOG.error(e);
            throw new IllegalStateException(e.getMessage());
         }
      }

      /**
       * create instances of WriteHandler threads for sending data.
       *
       * @param workerThreads is the number of threads to create.
       */
      private static WriteHandler[] create(int workerThreads, ThreadGroup tg, LinkedList backGroundThreads)
      {
         WriteHandler[] handlers = new WriteHandler[workerThreads];
         for (int looper = 0; looper < workerThreads; looper++)
         {
            handlers[looper] = new WriteHandler();

            Thread thread = new Thread(tg, handlers[looper], "nioWriteHandlerThread");
            thread.setDaemon(true);
            thread.start();
            backGroundThreads.add(thread);
         }
         return handlers;
      }

      /**
       * Add a new channel to be handled.
       *
       * @param channel
       */
      private SelectorWriteHandler add(SocketChannel channel)
      {
          return new SelectorWriteHandler(channel, SELECTOR, m_headerBuffer);
      }

      /**
       * Writes buffer to the specified socket connection.  This is always performed asynchronously.  If you want
       * to perform a synchrounous write, call notification.`get() which will block until the write operation is complete.
       * Best practice is to call notification.getException() which may return any exceptions that occured during the write
       * operation.
       *
       * @param channel      is where the buffer is written to.
       * @param buffer       is what we write.
       * @param notification may be specified if you want to know how many bytes were written and know if an exception