Examples with ServerStream io.grpc.internal.ServerStream used on opensource projects

Example 11 with ServerStream

use of io.grpc.internal.ServerStream in project grpc-java by grpc.

the class AbstractTransportTest method flowControlPushBack.

@Test
public void flowControlPushBack() throws Exception {
    server.start(serverListener);
    client = newClientTransport(server);
    runIfNotNull(client.start(mockClientTransportListener));
    MockServerTransportListener serverTransportListener = serverListener.takeListenerOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    serverTransport = serverTransportListener.transport;
    ClientStream clientStream = client.newStream(methodDescriptor, new Metadata());
    clientStream.start(mockClientStreamListener);
    StreamCreation serverStreamCreation = serverTransportListener.takeStreamOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);
    assertEquals(methodDescriptor.getFullMethodName(), serverStreamCreation.method);
    ServerStream serverStream = serverStreamCreation.stream;
    ServerStreamListener mockServerStreamListener = serverStreamCreation.listener;
    serverStream.writeHeaders(new Metadata());
    Answer<Void> closeStream = new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Exception {
            Object[] args = invocation.getArguments();
            ((InputStream) args[0]).close();
            return null;
        }
    };
    String largeMessage;
    {
        int size = 1 * 1024;
        StringBuffer sb = new StringBuffer(size);
        for (int i = 0; i < size; i++) {
            sb.append('a');
        }
        largeMessage = sb.toString();
    }
    doAnswer(closeStream).when(mockServerStreamListener).messageRead(any(InputStream.class));
    serverStream.request(1);
    verify(mockClientStreamListener, timeout(TIMEOUT_MS)).onReady();
    assertTrue(clientStream.isReady());
    final int maxToSend = 10 * 1024;
    int clientSent;
    // Verify that flow control will push back on client.
    for (clientSent = 0; clientStream.isReady(); clientSent++) {
        if (clientSent > maxToSend) {
            // It seems like flow control isn't working. _Surely_ flow control would have pushed-back
            // already. If this is normal, please configure the transport to buffer less.
            fail("Too many messages sent before isReady() returned false");
        }
        clientStream.writeMessage(methodDescriptor.streamRequest(largeMessage));
        clientStream.flush();
    }
    assertTrue(clientSent > 0);
    // Make sure there are at least a few messages buffered.
    for (; clientSent < 5; clientSent++) {
        clientStream.writeMessage(methodDescriptor.streamResponse(largeMessage));
        clientStream.flush();
    }
    doPingPong(serverListener);
    verify(mockServerStreamListener, timeout(TIMEOUT_MS)).messageRead(any(InputStream.class));
    doAnswer(closeStream).when(mockClientStreamListener).messageRead(any(InputStream.class));
    clientStream.request(1);
    verify(mockServerStreamListener, timeout(TIMEOUT_MS)).onReady();
    assertTrue(serverStream.isReady());
    int serverSent;
    // Verify that flow control will push back on server.
    for (serverSent = 0; serverStream.isReady(); serverSent++) {
        if (serverSent > maxToSend) {
            // It seems like flow control isn't working. _Surely_ flow control would have pushed-back
            // already. If this is normal, please configure the transport to buffer less.
            fail("Too many messages sent before isReady() returned false");
        }
        serverStream.writeMessage(methodDescriptor.streamResponse(largeMessage));
        serverStream.flush();
    }
    assertTrue(serverSent > 0);
    // Make sure there are at least a few messages buffered.
    for (; serverSent < 5; serverSent++) {
        serverStream.writeMessage(methodDescriptor.streamResponse(largeMessage));
        serverStream.flush();
    }
    doPingPong(serverListener);
    verify(mockClientStreamListener, timeout(TIMEOUT_MS)).messageRead(any(InputStream.class));
    serverStream.request(3);
    clientStream.request(3);
    doPingPong(serverListener);
    // times() is total number throughout the entire test
    verify(mockClientStreamListener, timeout(TIMEOUT_MS).times(4)).messageRead(any(InputStream.class));
    verify(mockServerStreamListener, timeout(TIMEOUT_MS).times(4)).messageRead(any(InputStream.class));
    // Request the rest
    serverStream.request(clientSent);
    clientStream.request(serverSent);
    verify(mockClientStreamListener, timeout(TIMEOUT_MS).times(clientSent)).messageRead(any(InputStream.class));
    verify(mockServerStreamListener, timeout(TIMEOUT_MS).times(serverSent)).messageRead(any(InputStream.class));
    verify(mockClientStreamListener, timeout(TIMEOUT_MS).times(2)).onReady();
    assertTrue(clientStream.isReady());
    verify(mockServerStreamListener, timeout(TIMEOUT_MS).times(2)).onReady();
    assertTrue(serverStream.isReady());
    // Request four more
    for (int i = 0; i < 5; i++) {
        clientStream.writeMessage(methodDescriptor.streamRequest(largeMessage));
        clientStream.flush();
        serverStream.writeMessage(methodDescriptor.streamResponse(largeMessage));
        serverStream.flush();
    }
    doPingPong(serverListener);
    verify(mockClientStreamListener, timeout(TIMEOUT_MS).times(clientSent + 4)).messageRead(any(InputStream.class));
    verify(mockServerStreamListener, timeout(TIMEOUT_MS).times(serverSent + 4)).messageRead(any(InputStream.class));
    // Drain exactly how many messages are left
    serverStream.request(1);
    clientStream.request(1);
    verify(mockServerStreamListener, timeout(TIMEOUT_MS).times(serverSent + 5)).messageRead(any(InputStream.class));
    verify(mockClientStreamListener, timeout(TIMEOUT_MS).times(clientSent + 5)).messageRead(any(InputStream.class));
    // And now check that the streams can still complete gracefully
    clientStream.writeMessage(methodDescriptor.streamRequest(largeMessage));
    clientStream.flush();
    clientStream.halfClose();
    doPingPong(serverListener);
    verify(mockServerStreamListener, never()).halfClosed();
    serverStream.request(1);
    verify(mockServerStreamListener, timeout(TIMEOUT_MS).times(serverSent + 6)).messageRead(any(InputStream.class));
    verify(mockServerStreamListener, timeout(TIMEOUT_MS)).halfClosed();
    serverStream.writeMessage(methodDescriptor.streamResponse(largeMessage));
    serverStream.flush();
    Status status = Status.OK.withDescription("... quite a lengthy discussion");
    serverStream.close(status, new Metadata());
    doPingPong(serverListener);
    verify(mockClientStreamListener, never()).closed(any(Status.class), any(Metadata.class));
    clientStream.request(1);
    verify(mockClientStreamListener, timeout(TIMEOUT_MS).times(clientSent + 6)).messageRead(any(InputStream.class));
    verify(mockServerStreamListener, timeout(TIMEOUT_MS)).closed(statusCaptor.capture());
    assertCodeEquals(Status.OK, statusCaptor.getValue());
    verify(mockClientStreamListener, timeout(TIMEOUT_MS)).closed(statusCaptor.capture(), any(Metadata.class));
    assertEquals(status.getCode(), statusCaptor.getValue().getCode());
    assertEquals(status.getDescription(), statusCaptor.getValue().getDescription());
}