public ResT receive(long timeoutMs) throws IOException {
    if (mCompleted) {
      return null;
    }
    if (mCanceled) {
      throw new CancelledException(formatErrorMessage("Stream is already canceled."));
    }

    long startMs = System.currentTimeMillis();
    while (true) {
      long waitedForMs = System.currentTimeMillis() - startMs;
      if (waitedForMs >= timeoutMs) {
        throw new DeadlineExceededException(formatErrorMessage(
            "Timeout waiting for response after %dms. clientClosed: %s clientCancelled: %s "
                + "serverClosed: %s", timeoutMs, mClosed, mCanceled, mClosedFromRemote));
      }

      // Wait for a minute max
      long waitMs = Math.min(timeoutMs - waitedForMs, Constants.MINUTE_MS);
      try {
        Object response = mResponses.poll(waitMs, TimeUnit.MILLISECONDS);
        if (response == null) {
          checkError(); // The stream could have errored while we were waiting
          // Log a warning before looping again
          LOG.warn("Client did not receive message from stream, will wait again. totalWaitMs: {} "
                  + "clientClosed: {} clientCancelled: {} serverClosed: {} description: {}",
              System.currentTimeMillis() - startMs, mClosed, mCanceled, mClosedFromRemote,
              mDescription);
          continue;
        }
        if (response == mResponseObserver) {
          mCompleted = true;
          return null;
        }
        checkError();
        return (ResT) response;
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        throw new CancelledException(
            formatErrorMessage("Interrupted while waiting for response."), e);
      }
    }
  }
@Test
  public void receiveErrorWhenBufferFull() throws Exception {
    WriteResponse[] responses = Stream.generate(() -> WriteResponse.newBuilder().build())
        .limit(BUFFER_SIZE).toArray(WriteResponse[]::new);
    for (WriteResponse response : responses) {
      mResponseObserver.onNext(response);
    }
    mResponseObserver.onError(Status.UNAUTHENTICATED.asRuntimeException());

    Exception e = assertThrows(UnauthenticatedException.class, () -> {
      for (WriteResponse response : responses) {
        WriteResponse actualResponse = mStream.receive(TIMEOUT);
        assertEquals(response, actualResponse);
      }
    });
    assertTrue(e.getMessage().contains(TEST_MESSAGE));
  }