@Test
 public void fourPeers() throws Exception {
   InboundMessageQueuer[] channels = {
     connectPeer(1), connectPeer(2), connectPeer(3), connectPeer(4)
   };
   Transaction tx = new Transaction(params);
   TransactionBroadcast broadcast = new TransactionBroadcast(peerGroup, tx);
   ListenableFuture<Transaction> future = broadcast.broadcast();
   assertFalse(future.isDone());
   // We expect two peers to receive a tx message, and at least one of the others must announce for
   // the future to
   // complete successfully.
   Message[] messages = {
     (Message) outbound(channels[0]),
     (Message) outbound(channels[1]),
     (Message) outbound(channels[2]),
     (Message) outbound(channels[3])
   };
   // 0 and 3 are randomly selected to receive the broadcast.
   assertEquals(tx, messages[0]);
   assertEquals(tx, messages[3]);
   assertNull(messages[1]);
   assertNull(messages[2]);
   Threading.waitForUserCode();
   assertFalse(future.isDone());
   inbound(channels[1], InventoryMessage.with(tx));
   pingAndWait(channels[1]);
   Threading.waitForUserCode();
   assertTrue(future.isDone());
 }
예제 #2
0
/**
 * A simple NIO MessageWriteTarget which handles all the business logic of a connection
 * (reading+writing bytes). Used only by the NioClient and NioServer classes
 */
class ConnectionHandler implements MessageWriteTarget {
  private static final org.slf4j.Logger log = LoggerFactory.getLogger(ConnectionHandler.class);

  private static final int BUFFER_SIZE_LOWER_BOUND = 4096;
  private static final int BUFFER_SIZE_UPPER_BOUND = 65536;

  private static final int OUTBOUND_BUFFER_BYTE_COUNT =
      Message.MAX_SIZE + 24; // 24 byte message header

  // We lock when touching local flags and when writing data, but NEVER when calling any methods
  // which leave this
  // class into non-Java classes.
  private final ReentrantLock lock = Threading.lock("nioConnectionHandler");

  @GuardedBy("lock")
  private final ByteBuffer readBuff;

  @GuardedBy("lock")
  private final SocketChannel channel;

  @GuardedBy("lock")
  private final SelectionKey key;

  @GuardedBy("lock")
  StreamParser parser;

  @GuardedBy("lock")
  private boolean closeCalled = false;

  @GuardedBy("lock")
  private long bytesToWriteRemaining = 0;

  @GuardedBy("lock")
  private final LinkedList<ByteBuffer> bytesToWrite = new LinkedList<ByteBuffer>();

  private Set<ConnectionHandler> connectedHandlers;

  public ConnectionHandler(StreamParserFactory parserFactory, SelectionKey key) throws IOException {
    this(
        parserFactory.getNewParser(
            ((SocketChannel) key.channel()).socket().getInetAddress(),
            ((SocketChannel) key.channel()).socket().getPort()),
        key);
    if (parser == null) throw new IOException("Parser factory.getNewParser returned null");
  }

  private ConnectionHandler(@Nullable StreamParser parser, SelectionKey key) {
    this.key = key;
    this.channel = checkNotNull(((SocketChannel) key.channel()));
    if (parser == null) {
      readBuff = null;
      closeConnection();
      return;
    }
    this.parser = parser;
    readBuff =
        ByteBuffer.allocateDirect(
            Math.min(
                Math.max(parser.getMaxMessageSize(), BUFFER_SIZE_LOWER_BOUND),
                BUFFER_SIZE_UPPER_BOUND));
    parser.setWriteTarget(this); // May callback into us (eg closeConnection() now)
    connectedHandlers = null;
  }

  public ConnectionHandler(
      StreamParser parser, SelectionKey key, Set<ConnectionHandler> connectedHandlers) {
    this(checkNotNull(parser), key);

    // closeConnection() may have already happened, in which case we shouldn't add ourselves to the
    // connectedHandlers set
    lock.lock();
    boolean alreadyClosed = false;
    try {
      alreadyClosed = closeCalled;
      this.connectedHandlers = connectedHandlers;
    } finally {
      lock.unlock();
    }
    if (!alreadyClosed) checkState(connectedHandlers.add(this));
  }

  @GuardedBy("lock")
  private void setWriteOps() {
    // Make sure we are registered to get updated when writing is available again
    key.interestOps(key.interestOps() | SelectionKey.OP_WRITE);
    // Refresh the selector to make sure it gets the new interestOps
    key.selector().wakeup();
  }

  // Tries to write any outstanding write bytes, runs in any thread (possibly unlocked)
  private void tryWriteBytes() throws IOException {
    lock.lock();
    try {
      // Iterate through the outbound ByteBuff queue, pushing as much as possible into the OS'
      // network buffer.
      Iterator<ByteBuffer> bytesIterator = bytesToWrite.iterator();
      while (bytesIterator.hasNext()) {
        ByteBuffer buff = bytesIterator.next();
        bytesToWriteRemaining -= channel.write(buff);
        if (!buff.hasRemaining()) bytesIterator.remove();
        else {
          setWriteOps();
          break;
        }
      }
      // If we are done writing, clear the OP_WRITE interestOps
      if (bytesToWrite.isEmpty()) key.interestOps(key.interestOps() & ~SelectionKey.OP_WRITE);
      // Don't bother waking up the selector here, since we're just removing an op, not adding
    } finally {
      lock.unlock();
    }
  }

  @Override
  public void writeBytes(byte[] message) throws IOException {
    lock.lock();
    try {
      // Network buffers are not unlimited (and are often smaller than some messages we may wish to
      // send), and
      // thus we have to buffer outbound messages sometimes. To do this, we use a queue of
      // ByteBuffers and just
      // append to it when we want to send a message. We then let tryWriteBytes() either send the
      // message or
      // register our SelectionKey to wakeup when we have free outbound buffer space available.

      if (bytesToWriteRemaining + message.length > OUTBOUND_BUFFER_BYTE_COUNT)
        throw new IOException("Outbound buffer overflowed");
      // Just dump the message onto the write buffer and call tryWriteBytes
      // TODO: Kill the needless message duplication when the write completes right away
      bytesToWrite.offer(ByteBuffer.wrap(Arrays.copyOf(message, message.length)));
      bytesToWriteRemaining += message.length;
      setWriteOps();
    } catch (IOException e) {
      lock.unlock();
      log.error("Error writing message to connection, closing connection", e);
      closeConnection();
      throw e;
    } catch (CancelledKeyException e) {
      lock.unlock();
      log.error("Error writing message to connection, closing connection", e);
      closeConnection();
      throw new IOException(e);
    }
    lock.unlock();
  }

  @Override
  // May NOT be called with lock held
  public void closeConnection() {
    try {
      channel.close();
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    connectionClosed();
  }

  private void connectionClosed() {
    boolean callClosed = false;
    lock.lock();
    try {
      callClosed = !closeCalled;
      closeCalled = true;
    } finally {
      lock.unlock();
    }
    if (callClosed) {
      checkState(connectedHandlers == null || connectedHandlers.remove(this));
      parser.connectionClosed();
    }
  }

  // Handle a SelectionKey which was selected
  // Runs unlocked as the caller is single-threaded (or if not, should enforce that handleKey is
  // only called
  // atomically for a given ConnectionHandler)
  public static void handleKey(SelectionKey key) {
    ConnectionHandler handler = ((ConnectionHandler) key.attachment());
    try {
      if (handler == null) return;
      if (!key.isValid()) {
        handler.closeConnection(); // Key has been cancelled, make sure the socket gets closed
        return;
      }
      if (key.isReadable()) {
        // Do a socket read and invoke the parser's receiveBytes message
        int read = handler.channel.read(handler.readBuff);
        if (read == 0) return; // Was probably waiting on a write
        else if (read == -1) { // Socket was closed
          key.cancel();
          handler.closeConnection();
          return;
        }
        // "flip" the buffer - setting the limit to the current position and setting position to 0
        handler.readBuff.flip();
        // Use parser.receiveBytes's return value as a check that it stopped reading at the right
        // location
        int bytesConsumed = checkNotNull(handler.parser).receiveBytes(handler.readBuff);
        checkState(handler.readBuff.position() == bytesConsumed);
        // Now drop the bytes which were read by compacting readBuff (resetting limit and keeping
        // relative
        // position)
        handler.readBuff.compact();
      }
      if (key.isWritable()) handler.tryWriteBytes();
    } catch (Exception e) {
      // This can happen eg if the channel closes while the thread is about to get killed
      // (ClosedByInterruptException), or if handler.parser.receiveBytes throws something
      log.error("Error handling SelectionKey: {}", Throwables.getRootCause(e).getMessage());
      handler.closeConnection();
    }
  }
}
  @Test
  public void peerGroupWalletIntegration() throws Exception {
    // Make sure we can create spends, and that they are announced. Then do the same with offline
    // mode.

    // Set up connections and block chain.
    VersionMessage ver = new VersionMessage(params, 2);
    ver.localServices = VersionMessage.NODE_NETWORK;
    InboundMessageQueuer p1 = connectPeer(1, ver);
    InboundMessageQueuer p2 = connectPeer(2);

    // Send ourselves a bit of money.
    Block b1 = TestUtils.makeSolvedTestBlock(blockStore, address);
    inbound(p1, b1);
    pingAndWait(p1);
    assertNull(outbound(p1));
    assertEquals(Utils.toNanoCoins(50, 0), wallet.getBalance());

    // Check that the wallet informs us of changes in confidence as the transaction ripples across
    // the network.
    final Transaction[] transactions = new Transaction[1];
    wallet.addEventListener(
        new AbstractWalletEventListener() {
          @Override
          public void onTransactionConfidenceChanged(Wallet wallet, Transaction tx) {
            transactions[0] = tx;
          }
        });

    // Now create a spend, and expect the announcement on p1.
    Address dest = new ECKey().toAddress(params);
    Wallet.SendResult sendResult = wallet.sendCoins(peerGroup, dest, Utils.toNanoCoins(1, 0));
    assertNotNull(sendResult.tx);
    Threading.waitForUserCode();
    assertFalse(sendResult.broadcastComplete.isDone());
    assertEquals(transactions[0], sendResult.tx);
    assertEquals(0, transactions[0].getConfidence().numBroadcastPeers());
    transactions[0] = null;
    Transaction t1 = (Transaction) outbound(p1);
    assertNotNull(t1);
    // 49 BTC in change.
    assertEquals(Utils.toNanoCoins(49, 0), t1.getValueSentToMe(wallet));
    // The future won't complete until it's heard back from the network on p2.
    InventoryMessage inv = new InventoryMessage(params);
    inv.addTransaction(t1);
    inbound(p2, inv);
    pingAndWait(p2);
    Threading.waitForUserCode();
    assertTrue(sendResult.broadcastComplete.isDone());
    assertEquals(transactions[0], sendResult.tx);
    assertEquals(1, transactions[0].getConfidence().numBroadcastPeers());
    // Confirm it.
    Block b2 = TestUtils.createFakeBlock(blockStore, t1).block;
    inbound(p1, b2);
    pingAndWait(p1);
    assertNull(outbound(p1));

    // Do the same thing with an offline transaction.
    peerGroup.removeWallet(wallet);
    Wallet.SendRequest req = Wallet.SendRequest.to(dest, Utils.toNanoCoins(2, 0));
    req.ensureMinRequiredFee = false;
    Transaction t3 = checkNotNull(wallet.sendCoinsOffline(req));
    assertNull(outbound(p1)); // Nothing sent.
    // Add the wallet to the peer group (simulate initialization). Transactions should be announced.
    peerGroup.addWallet(wallet);
    // Transaction announced to the first peer.
    assertEquals(t3.getHash(), ((Transaction) outbound(p1)).getHash());
  }