org.apache.thrift.transport.TNonblockingTransport Java Examples

@Override
protected void processKey(SelectionKey key) throws IOException
{
    if (!key.isAcceptable())
        return;

    try
    {
        // accept the connection
        SelectorThread selector = selectorLoadBalancer.nextSelector();
        selector.subscribe((TNonblockingTransport) serverTransport.accept());
        selector.wakeupSelector();
    }
    catch (TTransportException tte)
    {
        // accept() shouldn't be NULL if fine because are are raising for a socket
        logger.debug("Non-fatal exception trying to accept!", tte);
    }
}
 

/**
 * Accept a new connection.
 */
private void handleAccept() {
  final TNonblockingTransport client = doAccept();
  if (client != null) {
    // Pass this connection to a selector thread
    final SelectorThread targetThread = threadChooser.nextThread();

    if (args.acceptPolicy == Args.AcceptPolicy.FAST_ACCEPT || invoker == null) {
      doAddAccept(targetThread, client);
    } else {
      // FAIR_ACCEPT
      try {
        invoker.submit(new Runnable() {
          public void run() {
            doAddAccept(targetThread, client);
          }
        });
      } catch (RejectedExecutionException rx) {
        LOGGER.warn("ExecutorService rejected accept registration!", rx);
        // close immediately
        client.close();
      }
    }
  }
}
 

/**
 * Hands off an accepted connection to be handled by this thread. This
 * method will block if the queue for new connections is at capacity.
 *
 * @param accepted
 *          The connection that has been accepted.
 * @return true if the connection has been successfully added.
 */
public boolean addAcceptedConnection(TNonblockingTransport accepted) {
  try {
    while (!acceptedQueue.offer(accepted, 200, TimeUnit.MILLISECONDS)) {
      // If server is stopped, then return false.
      if (stopped_) {
        return false;
      }
    }
  } catch (InterruptedException e) {
    LOGGER.warn("Interrupted while adding accepted connection!", e);
    return false;
  }
  selector.wakeup();
  return true;
}
 

@Test
public void async() throws Exception {
  TestUtil.setGlobalTracer(tracer);

  startAsyncServer();

  final Factory protocolFactory = new Factory();
  final TNonblockingTransport transport = new TNonblockingSocket("localhost", port);
  final TAsyncClientManager clientManager = new TAsyncClientManager();
  final AsyncClient asyncClient = new AsyncClient(protocolFactory, clientManager, transport);
  final AtomicInteger counter = new AtomicInteger();
  asyncClient.say("Async", "World", new AsyncMethodCallback<String>() {
    @Override
    public void onComplete(final String response) {
      assertEquals("Say Async World", response);
      assertNotNull(GlobalTracer.get().activeSpan());
      counter.incrementAndGet();
    }

    @Override
    public void onError(final Exception exception) {
      exception.printStackTrace();
    }
  });

  await().atMost(15, TimeUnit.SECONDS).until(TestUtil.reportedSpansSize(tracer), equalTo(2));
  assertEquals(1, counter.get());

  final List<MockSpan> spans = tracer.finishedSpans();
  assertEquals(2, spans.size());

  assertNull(tracer.activeSpan());
  verify(tracer, times(2)).buildSpan(anyString());
  verify(tracer, times(1)).inject(any(SpanContext.class), any(Format.class), any());
}
 

public int writeTo(TNonblockingTransport transport, int start, int count) throws IOException
{
    ByteBuffer dup = buffer.duplicate();
    dup.position(start).limit(start + count);

    return transport.write(dup);
}
 

@Override
protected void selectorIterationComplete() throws IOException
{
    TNonblockingTransport newClient;

    while ((newClient = newConnections.poll()) != null)
    {
        SelectionKey clientKey = newClient.registerSelector(selector, SelectionKey.OP_READ);
        clientKey.attach(new Message(newClient, clientKey, thriftFactories, useHeapBasedAllocation, alwaysReallocateBuffers));
    }
}
 

public Message(TNonblockingTransport trans,
               SelectionKey key,
               ThriftFactories factories,
               boolean heapBasedAllocation,
               boolean reallocateBuffers)
{
    frameSizeBuffer = Buffer.allocate(4, heapBasedAllocation);
    transport = trans;
    selectionKey = key;
    thriftFactories = factories;
    useHeapBasedAllocation = heapBasedAllocation;
    alwaysReallocateBuffers = reallocateBuffers;
}
 

protected final Socket getRootSocket(Object target) {
    if (target instanceof TNonblockingTransportFieldGetter) {
        TNonblockingTransport inTrans = ((TNonblockingTransportFieldGetter) target)._$PINPOINT$_getTNonblockingTransport();
        if (inTrans != null) {
            if (inTrans instanceof SocketFieldAccessor) {
                return ((SocketFieldAccessor) inTrans)._$PINPOINT$_getSocket();
            } else {
                if (isDebug) {
                    logger.debug("Invalid target object. Need field accessor({}).", SocketFieldAccessor.class.getName());
                }
            }
        }
    }
    return null;
}
 

@Override
public void after(Object target, Object[] args, Object result, Throwable throwable) {
    if (isDebug) {
        logger.afterInterceptor(target, args, result, throwable);
    }
    if (validate(target)) {
        TNonblockingTransport transport = ((TNonblockingTransportFieldGetter)target)._$PINPOINT$_getTNonblockingTransport();
        if (validateTransport(transport)) {
            SocketAddress socketAddress = ((SocketAddressFieldAccessor)transport)._$PINPOINT$_getSocketAddress();
            ((SocketAddressFieldAccessor)target)._$PINPOINT$_setSocketAddress(socketAddress);
        }
    }
}
 

/**
 * 获取客户端AsyncService对象
 */
@Override
public AsyncService getAsyncService(TNonblockingTransport transport, String serviceName) throws IkasoaException {
	if (ObjectUtil.isNull(transport))
		throw new IllegalArgumentException("'transport' is null !");
	try {
		return StringUtil.isEmpty(serviceName)
				? new AsyncServiceClientImpl((TProtocolFactory) new TCompactProtocol.Factory(), transport)
				: new AsyncServiceClientImpl(new AsyncMultiplexedProtocolFactory(serviceName), transport);
	} catch (IOException e) {
		throw new IkasoaException(e);
	}
}
 

private void registerAccepted(TNonblockingTransport accepted) {
  SelectionKey clientKey = null;
  try {
    clientKey = accepted.registerSelector(selector, SelectionKey.OP_READ);

    FrameBuffer frameBuffer = new FrameBuffer(accepted, clientKey, SelectorThread.this);
    clientKey.attach(frameBuffer);
  } catch (IOException e) {
    LOGGER.warn("Failed to register accepted connection to selector!", e);
    if (clientKey != null) {
      cleanupSelectionKey(clientKey);
    }
    accepted.close();
  }
}
 

private void processAcceptedConnections() {
  // Register accepted connections
  while (!stopped_) {
    TNonblockingTransport accepted = acceptedQueue.poll();
    if (accepted == null) {
      break;
    }
    registerAccepted(accepted);
  }
}
 

@Test
public void testAsyncServerRule() throws Exception {
  String responseMsg = "asyncHello";
  Response expectHelloResponse = new Response(200, responseMsg);
  server.setExpectReturn("sayHello", expectHelloResponse);

  TNonblockingTransport transport = new TNonblockingSocket("127.0.0.1",
                                                           9999,
                                                           3000);
  HelloService.AsyncIface client = new HelloService.AsyncClient(new TBinaryProtocol.Factory(),
                                                                new TAsyncClientManager(),
                                                                transport);
  Request request = new Request();
  request.setMsg("hello async!");
  Response result = new Response();
  Stopwatch stopwatch = Stopwatch.createStarted();
  client.sayHello(request, new AsyncMethodCallback() {
    @Override
    public void onComplete(Object response) {
      System.out.println("cost:" + stopwatch.elapsed(TimeUnit.MILLISECONDS));
      result.setCode(((Response)response).getCode());
      result.setResponseMsg(((Response)response).getResponseMsg());
    }

    @Override
    public void onError(Exception exception) {
      System.out.println(exception.getMessage());
    }
  });
  Thread.sleep(50);
  System.out.println("final result:" + result);
  Assert.assertTrue(result.getCode() == 200);
  Assert.assertTrue(responseMsg.equals(result.getResponseMsg()));
}
 

@Test
public void asyncWithoutArgs() throws Exception {
  TestUtil.setGlobalTracer(tracer);

  startAsyncServer();

  final Factory protocolFactory = new Factory();

  final TNonblockingTransport transport = new TNonblockingSocket("localhost", port);
  final TAsyncClientManager clientManager = new TAsyncClientManager();
  final AsyncClient asyncClient = new AsyncClient(protocolFactory, clientManager, transport);
  final AtomicInteger counter = new AtomicInteger();
  asyncClient.withoutArgs(new AsyncMethodCallback<String>() {
    @Override
    public void onComplete(final String response) {
      assertEquals("no args", response);
      assertNotNull(GlobalTracer.get().activeSpan());
      counter.incrementAndGet();
    }

    @Override
    public void onError(final Exception exception) {
      exception.printStackTrace();
    }
  });

  await().atMost(15, TimeUnit.SECONDS).until(TestUtil.reportedSpansSize(tracer), equalTo(2));
  assertEquals(1, counter.get());

  final List<MockSpan> spans = tracer.finishedSpans();
  assertEquals(2, spans.size());

  assertNull(tracer.activeSpan());
  verify(tracer, times(2)).buildSpan(anyString());
  verify(tracer, times(1)).inject(any(SpanContext.class), any(Format.class), any());
}
 

@Test
public void asyncMany() throws Exception {
  TestUtil.setGlobalTracer(tracer);

  startAsyncServer();

  final AtomicInteger counter = new AtomicInteger();
  for (int i = 0; i < 4; ++i) {
    final Factory protocolFactory = new Factory();
    final TNonblockingTransport transport = new TNonblockingSocket("localhost", port);
    final TAsyncClientManager clientManager = new TAsyncClientManager();
    final AsyncClient asyncClient = new AsyncClient(protocolFactory, clientManager, transport);
    asyncClient.withDelay(1, new AsyncMethodCallback<String>() {
      @Override
      public void onComplete(final String response) {
        assertEquals("delay 1", response);
        assertNotNull(GlobalTracer.get().activeSpan());
        counter.incrementAndGet();
      }

      @Override
      public void onError(final Exception exception) {
        exception.printStackTrace();
      }
    });
  }

  await().atMost(15, TimeUnit.SECONDS).until(TestUtil.reportedSpansSize(tracer), equalTo(8));
  assertEquals(4, counter.get());

  final List<MockSpan> spans = tracer.finishedSpans();
  assertEquals(8, spans.size());

  assertNull(tracer.activeSpan());
  verify(tracer, times(8)).buildSpan(anyString());
  verify(tracer, times(4)).inject(any(SpanContext.class), any(Format.class), any());
}
 

@Test
public void oneWayAsync() throws Exception {
  TestUtil.setGlobalTracer(tracer);

  startAsyncServer();

  final Factory protocolFactory = new Factory();
  final TNonblockingTransport transport = new TNonblockingSocket("localhost", port);
  final TAsyncClientManager clientManager = new TAsyncClientManager();
  final AsyncClient asyncClient = new AsyncClient(protocolFactory, clientManager, transport);
  final AtomicInteger counter = new AtomicInteger();
  asyncClient.oneWay(new AsyncMethodCallback<Void>() {
    @Override
    public void onComplete(final Void response) {
      assertNotNull(GlobalTracer.get().activeSpan());
      counter.incrementAndGet();
    }

    @Override
    public void onError(final Exception exception) {
      exception.printStackTrace();
    }
  });

  await().atMost(15, TimeUnit.SECONDS).until(TestUtil.reportedSpansSize(tracer), equalTo(2));
  assertEquals(1, counter.get());

  final List<MockSpan> spans = tracer.finishedSpans();
  assertEquals(2, spans.size());

  assertNull(tracer.activeSpan());
  verify(tracer, times(2)).buildSpan(anyString());
  verify(tracer, times(1)).inject(any(SpanContext.class), any(Format.class), any());
}
 

private static BlockingQueue<TNonblockingTransport> createDefaultAcceptQueue(int queueSize) {
  if (queueSize == 0) {
    // Unbounded queue
    return new LinkedBlockingQueue<TNonblockingTransport>();
  }
  return new ArrayBlockingQueue<TNonblockingTransport>(queueSize);
}
 

private TNonblockingTransport doAccept() {
  try {
    return (TNonblockingTransport) serverTransport.accept();
  } catch (TTransportException tte) {
    // something went wrong accepting.
    LOGGER.warn("Exception trying to accept!", tte);
    return null;
  }
}
 

public TNonblockingTransport getTransport() {
    return this.trans_;
}
 

public int readFrom(TNonblockingTransport transport) throws IOException
{
    return transport.read(buffer);
}
 

public void subscribe(TNonblockingTransport newClient)
{
    newConnections.add(newClient); // always returns true as queue is unbounded
}
 

@Test
public void test_AsyncClient() throws Throwable {
	Random rnd = new Random(System.nanoTime());

	TProtocolFactory[] protfacs = new TProtocolFactory[] { new TCompactProtocol.Factory(),
			new TBinaryProtocol.Factory(), new TJSONProtocol.Factory(),
			new TSimpleJSONProtocol.Factory(TCalculator.Iface.class, false) };

	TProtocolFactory protocolFactory = protfacs[rnd.nextInt(protfacs.length)];

	System.out.println("protocolFactory: " + protocolFactory);

	TAsyncClientManager clientManager = new TAsyncClientManager();
	TNonblockingTransport transport = new TNonblockingSocket(HOST, PORT);
	TCalculator.AsyncClient client = new TCalculator.AsyncClient(protocolFactory, clientManager, transport);
	final int num1 = rnd.nextInt(Integer.MAX_VALUE / 2 - 1);
	final int num2 = rnd.nextInt(Integer.MAX_VALUE / 2 - 1);

	final CountDownLatch latch = new CountDownLatch(1);
	final Throwable[] exceptions = new Throwable[1];
	AsyncMethodCallback<TCalculator.AsyncClient.add_call> resultHandler = new AsyncMethodCallback<TCalculator.AsyncClient.add_call>() {
		@Override
		public void onComplete(TCalculator.AsyncClient.add_call response) {
			System.out.println("onComplete!");
			try {
				int result = response.getResult();
				Assert.assertEquals(num1 + num2, result);
			} catch (Throwable e) {
				exceptions[0] = e;
			} finally {
				latch.countDown();
			}
		}

		@Override
		public void onError(Exception exception) {
			System.err.println("onError!");
			exception.printStackTrace();
			latch.countDown();
		}

	};
	client.add(num1, num2, resultHandler);
	latch.await();
	transport.close();
	if (exceptions[0] != null) {
		throw exceptions[0];
	}
}
 

public CustomFrameBuffer(final TNonblockingTransport trans,
       final SelectionKey selectionKey,
       final AbstractSelectThread selectThread) {
super(trans, selectionKey, selectThread);
     }
 

/**
 * 获取客户端AsyncService对象
 */
@Override
public AsyncService getAsyncService(TNonblockingTransport transport) throws IkasoaException {
	return getAsyncService(transport, null);
}
 

public CallBack(String arg, AsyncMethodCallback<String> resultHandler, TAsyncClient client,
		TProtocolFactory protocolFactory, TNonblockingTransport transport) throws TException {
	super(client, protocolFactory, transport, resultHandler, false);
	this.arg = arg;
}
 

public AsyncServiceClientImpl(TProtocolFactory protocolFactory, TNonblockingTransport transport)
		throws IOException {
	super(protocolFactory, new TAsyncClientManager(), transport);
}
 

public AsyncServiceClientImpl(TProtocolFactory protocolFactory, TAsyncClientManager manager,
		TNonblockingTransport transport) {
	super(protocolFactory, manager, transport);
}
 

private void doAddAccept(SelectorThread thread, TNonblockingTransport client) {
  if (!thread.addAcceptedConnection(client)) {
    client.close();
  }
}
 

/**
 * Start writing accumulated content to the given transport
 * or continue from the previous stop point (this method is stateful).
 *
 * @param transport The transport to write contents into.
 *
 * @return The number of bytes actually written to the transport.
 *
 * @throws IOException on any I/O related error, such as socket disconnect.
 */
public int streamTo(TNonblockingTransport transport) throws IOException
{
    int bytesWritten = buf.writeTo(transport, streamPos, count - streamPos);
    streamPos += bytesWritten;
    return bytesWritten;
}
 

/**
 * Set up the SelectorThread with a specified queue for connections.
 *
 * @param acceptedQueue
 *          The BlockingQueue implementation for holding incoming accepted
 *          connections.
 * @throws IOException
 *           if a selector cannot be created.
 */
public SelectorThread(BlockingQueue<TNonblockingTransport> acceptedQueue) throws IOException {
  this.acceptedQueue = acceptedQueue;
}