org.apache.thrift.transport.TNonblockingSocket Java Examples

@Test
public void testAysncMultiplexedServiceImpl() {
	int serverPort = ServerUtil.getNewPort();
	Map<String, TProcessor> processorMap = MapUtil.newHashMap();
	processorMap.put("testAysncService1", new ServiceProcessor(new TestThriftServiceImpl1()));
	processorMap.put("testAysncService2", new ServiceProcessor(new TestThriftServiceImpl2()));
	MultiplexedProcessor p = new MultiplexedProcessor(processorMap);
	thriftServerConfiguration.setProcessorFactory(new TProcessorFactory(p));
	Factory factory = new GeneralFactory(thriftServerConfiguration);
	ThriftServer thriftServer = factory.getThriftServer("testAysncMultiplexedService", serverPort, p);
	thriftServer.run();
	waiting();
	try {
		AsyncService service1 = factory
				.getAsyncService(new TNonblockingSocket(TestConstants.LOCAL_HOST, serverPort), "testAysncService1");
		service1.get(testString1, new TestCallback1());
		AsyncService service2 = factory
				.getAsyncService(new TNonblockingSocket(TestConstants.LOCAL_HOST, serverPort), "testAysncService2");
		service2.get(testString2, new TestCallback2());
		waiting();
	} catch (Exception e) {
		fail();
	} finally {
		thriftServer.stop();
	}
}
 

protected TNonblockingSocket acceptImpl() throws TTransportException {
    if (serverSocket_ == null) {
        throw new TTransportException(TTransportException.NOT_OPEN, "No underlying server socket.");
    }
    try {
        SocketChannel socketChannel = serverSocketChannel.accept();
        if (socketChannel == null) {
            return null;
        }

        TNonblockingSocket tsocket = new TNonblockingSocket(socketChannel);
        tsocket.setTimeout(0); // disabling client timeout
        tsocket.getSocketChannel().socket().setKeepAlive(true);
        tsocket.getSocketChannel().socket().setSendBufferSize(config.getSocketSendBufferBytes());
        tsocket.getSocketChannel().socket().setReceiveBufferSize(config.getSocketRecvBufferBytes());
        return tsocket;
    } catch (IOException iox) {
        throw new TTransportException(iox);
    }
}
 

private boolean validate(Object target, Object[] args) {
    if (!(target instanceof TNonblockingSocket)) {
        return false;
    }
    if (args.length != 3) {
        return false;
    }
    if (!(target instanceof SocketFieldAccessor)) {
        if (isDebug) {
            logger.debug("Invalid target object. Need field accessor({}).", SocketFieldAccessor.class.getName());
        }
        return false;
    }
    if (!(target instanceof SocketAddressFieldAccessor)) {
        if (isDebug) {
            logger.debug("Invalid target object. Need field accessor({}).", SocketAddressFieldAccessor.class.getName());
        }
        return false;
    }
    return true;
}
 

@Test
public void testAysncServiceImpl() {
	int serverPort = ServerUtil.getNewPort();
	TProcessor p = new ServiceProcessor(new TestThriftServiceImpl1());
	thriftServerConfiguration.setProcessorFactory(new TProcessorFactory(p));
	Factory factory = new GeneralFactory(thriftServerConfiguration);
	ThriftServer thriftServer = factory.getThriftServer(serverPort, new TestThriftServiceImpl1());
	thriftServer.run();
	waiting();
	try {
		AsyncService service = factory
				.getAsyncService(new TNonblockingSocket(TestConstants.LOCAL_HOST, serverPort));
		service.get(testString1, new TestCallback1());
		waiting();
	} catch (Exception e) {
		fail();
	} finally {
		thriftServer.stop();
	}
}
 

@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());
}
 

public ThriftContext(Runnable invocation) {
    //command should be instanceof Invocation
    try {
        AbstractNonblockingServer.FrameBuffer fb = (AbstractNonblockingServer.FrameBuffer) frameBufferField.get(invocation);
        TNonblockingSocket socket = (TNonblockingSocket) transField.get(fb);
        InetSocketAddress address = (InetSocketAddress) socket.getSocketChannel().getRemoteAddress();
        remoteAddress = address.getAddress().getHostAddress();
        LOGGER.debug("get thrift client address:{}", remoteAddress);
    } catch (Exception e) {
        LOGGER.error("create thrift context error", e);
    }
}
 

@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 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 testAysncThriftServerImpl() {
	int serverPort = ServerUtil.getNewPort();
	ThriftServerConfiguration thriftServerConfiguration = new ThriftServerConfiguration();
	thriftServerConfiguration.setProtocolFactory(new TCompactProtocol.Factory());
	thriftServerConfiguration.setProcessorFactory(new TProcessorFactory(processor));
	thriftServerConfiguration.setServerArgsAspect(new ServerArgsAspect() {
		@Override
		public TThreadPoolServer.Args tThreadPoolServerArgsAspect(TThreadPoolServer.Args args) {
			args.stopTimeoutVal = 1;
			return args;
		}
	});
	Factory factory = new GeneralFactory(thriftServerConfiguration);
	ThriftServer thriftServer = factory.getThriftServer(serverName, serverPort, processor);
	thriftServer.run();
	waiting();
	try {
		ThriftSimpleService.AsyncClient thriftClient = new ThriftSimpleService.AsyncClient(
				new TCompactProtocol.Factory(), new TAsyncClientManager(),
				new TNonblockingSocket(TestConstants.LOCAL_HOST, serverPort));
		thriftClient.get(TestConstants.TEST_STRING, new TestCallback());
		waiting();
	} catch (Exception e) {
		fail();
	} finally {
		thriftServer.stop();
	}
}
 

@Override
protected boolean requestInvoke(FrameBuffer frameBuffer)
{
    TNonblockingSocket socket = (TNonblockingSocket)((CustomFrameBuffer)frameBuffer).getTransport();
    ThriftSessionManager.instance.setCurrentSocket(socket.getSocketChannel().socket().getRemoteSocketAddress());
    frameBuffer.invoke();
    return true;
}
 

@Override
public void after(Object target, Object[] args, Object result, Throwable throwable) {
    if (isDebug) {
        logger.afterInterceptor(target, args, result, throwable);
    }
    if (validate(target, args)) {
        Socket socket = ((TNonblockingSocket)target).getSocketChannel().socket();
        ((SocketFieldAccessor)target)._$PINPOINT$_setSocket(socket);
        if (args[2] instanceof SocketAddress) {
            SocketAddress socketAddress = (SocketAddress)args[2];
            ((SocketAddressFieldAccessor)target)._$PINPOINT$_setSocketAddress(socketAddress);
        }
    }
}
 

@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 ThriftContext(Runnable invocation) {
    //command should be instanceof Invocation
    try {
        AbstractNonblockingServer.FrameBuffer fb = (AbstractNonblockingServer.FrameBuffer) frameBufferField.get(invocation);
        TNonblockingSocket socket = (TNonblockingSocket) transField.get(fb);
        InetSocketAddress address = (InetSocketAddress) socket.getSocketChannel().getRemoteAddress();
        remoteAddress = address.getAddress().getHostAddress();
        LOGGER.debug("get thrift client address:{}", remoteAddress);
    } catch (Exception e) {
        LOGGER.error("create thrift context error", e);
    }
}
 

@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 static void main(String[] args) 
        throws IOException, InterruptedException, TException {
    //Async client and I/O stack setup
    TNonblockingSocket trans_ep = new TNonblockingSocket("localhost", 9090);
    TAsyncClientManager client_man = new TAsyncClientManager();
    TradeReporting.TradeHistory.AsyncClient client = 
        new TradeReporting.TradeHistory.AsyncClient(new TBinaryProtocol.Factory(),
                                                    client_man, trans_ep);

    //get_last_sale() async callback handler
    WaitableCallback<TradeReport> wc = 
            new WaitableCallback<TradeReport>() {

        @Override
        public void onComplete(TradeReport tr) {
            try {
                System.out.println("[Client] received [" + tr.seq_num + "] " + 
                                   tr.symbol + " : " + tr.size + " @ " + tr.price);
            } finally {
                complete();
            }
        }
    };

    //Make async calls
    wc.reset();
    client.get_last_sale("IBM", wc);
    System.out.println("[Client] get_last_sale() executing asynch...");
    wc.wait(500);
    wc.reset();
    client.get_last_sale("F", wc);
    wc.wait(25000);

    //Make an async call which will time out
    client.setTimeout(1000);
    wc.reset();
    client.get_last_sale("GE", wc);
    wc.wait(5000);

    //Shutdown async client manager and close network socket
    client_man.stop();
    trans_ep.close();
}
 

@Override
protected void beforeInvoke(Message buffer)
{
    TNonblockingSocket socket = (TNonblockingSocket) buffer.transport;
    ThriftSessionManager.instance.setCurrentSocket(socket.getSocketChannel().socket().getRemoteSocketAddress());
}
 

public void beforeClose(Message buffer)
{
    TNonblockingSocket socket = (TNonblockingSocket) buffer.transport;
    ThriftSessionManager.instance.connectionComplete(socket.getSocketChannel().socket().getRemoteSocketAddress());
}
 

private AsyncEchoTestClient(TestEnvironment environment) throws IOException {
    this.environment = environment;
    this.transport = new TNonblockingSocket(this.environment.getServerIp(), this.environment.getPort());
    this.asyncClient = new EchoService.AsyncClient(this.environment.getProtocolFactory(), this.asyncClientManager, this.transport);
}