Examples with StreamObserver org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver used on opensource projects

Example 36 with StreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver in project alluxio by Alluxio.

the class AlluxioBlockStoreTest method getOutStreamWithReplicated.

@Test
public void getOutStreamWithReplicated() throws Exception {
    File file = File.createTempFile("test", ".tmp");
    CreateLocalBlockResponse response = CreateLocalBlockResponse.newBuilder().setPath(file.getAbsolutePath()).build();
    when(mWorkerClient.createLocalBlock(any(StreamObserver.class))).thenAnswer(new Answer() {

        public Object answer(InvocationOnMock invocation) {
            StreamObserver<CreateLocalBlockResponse> observer = invocation.getArgument(0, StreamObserver.class);
            observer.onNext(response);
            return mStreamObserver;
        }
    });
    when(mContext.getCachedWorkers()).thenReturn(Lists.newArrayList(new BlockWorkerInfo(WORKER_NET_ADDRESS_LOCAL, -1, -1), new BlockWorkerInfo(WORKER_NET_ADDRESS_REMOTE, -1, -1)));
    OutStreamOptions options = OutStreamOptions.defaults(mClientContext).setBlockSizeBytes(BLOCK_LENGTH).setLocationPolicy(new MockBlockLocationPolicy(Lists.newArrayList(WORKER_NET_ADDRESS_LOCAL, WORKER_NET_ADDRESS_REMOTE))).setWriteType(WriteType.MUST_CACHE).setReplicationMin(2);
    BlockOutStream stream = mBlockStore.getOutStream(BLOCK_ID, BLOCK_LENGTH, options);
    assertEquals(alluxio.client.block.stream.BlockOutStream.class, stream.getClass());
}

Example 37 with StreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver in project alluxio by Alluxio.

the class GrpcDataReaderTest method validateReadRequestSent.

/**
 * Validates the read request sent.
 *
 * @param client the worker client
 * @param offset the offset
 * @param length the length
 */
private void validateReadRequestSent(final BlockWorkerClient client, long offset, long length, boolean closed, int chunkSize) throws TimeoutException, InterruptedException {
    ArgumentCaptor<ReadRequest> requestCaptor = ArgumentCaptor.forClass(ReadRequest.class);
    verify(mRequestObserver, atLeastOnce()).onNext(requestCaptor.capture());
    ArgumentCaptor<StreamObserver> captor = ArgumentCaptor.forClass(StreamObserver.class);
    verify(mClient).readBlock(captor.capture());
    List<ReadRequest> readRequests = requestCaptor.getAllValues();
    assertTrue(!readRequests.isEmpty());
    captor.getValue().onCompleted();
    long lastOffset = offset;
    for (int i = 0; i < readRequests.size(); i++) {
        ReadRequest readRequest = readRequests.get(i);
        if (i == 0) {
            assertTrue(readRequest != null);
            assertEquals(BLOCK_ID, readRequest.getBlockId());
            assertEquals(offset, readRequest.getOffset());
            assertEquals(length, readRequest.getLength());
            assertEquals(chunkSize, readRequest.getChunkSize());
        } else {
            assertTrue(readRequest.hasOffsetReceived());
            assertTrue(readRequest.getOffsetReceived() > lastOffset);
            assertTrue(readRequest.getOffsetReceived() <= length);
            lastOffset = readRequest.getOffsetReceived();
        }
    }
    verify(mRequestObserver, closed ? atLeastOnce() : never()).onCompleted();
}

Example 38 with StreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver in project alluxio by Alluxio.

the class GrpcDataReaderTest method setReadResponses.

/**
 * Sets read responses which will be returned when readBlock() is invoked.
 *
 * @param client the worker client
 * @param length the length
 * @param start the start position to calculate the checksum
 * @param end the end position to calculate the checksum
 * @return the checksum
 */
private long setReadResponses(final BlockWorkerClient client, final long length, final long start, final long end) {
    long checksum = 0;
    long pos = 0;
    long remaining = length;
    ArgumentCaptor<StreamObserver> captor = ArgumentCaptor.forClass(StreamObserver.class);
    verify(mClient).readBlock(captor.capture());
    StreamObserver<ReadResponse> responseObserver = captor.getValue();
    List<ReadResponse> responses = new ArrayList<>();
    while (remaining > 0) {
        int bytesToSend = (int) Math.min(remaining, CHUNK_SIZE);
        byte[] data = new byte[bytesToSend];
        RANDOM.nextBytes(data);
        responses.add(ReadResponse.newBuilder().setChunk(Chunk.newBuilder().setData(ByteString.copyFrom(data))).build());
        remaining -= bytesToSend;
        for (int i = 0; i < data.length; i++) {
            if (pos >= start && pos <= end) {
                checksum += BufferUtils.byteToInt(data[i]);
            }
            pos++;
        }
    }
    EXECUTOR.submit(() -> {
        for (ReadResponse response : responses) {
            responseObserver.onNext(response);
        }
        responseObserver.onCompleted();
    });
    return checksum;
}

Example 39 with StreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver in project alluxio by Alluxio.

the class BlockMasterWorkerServiceHandlerTest method registerLeaseExpired.

@Test
public void registerLeaseExpired() {
    long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
    // Prepare LocationBlockIdListEntry objects
    BlockStoreLocation loc = new BlockStoreLocation("MEM", 0);
    BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder().setTierAlias(loc.tierAlias()).setMediumType(loc.mediumType()).build();
    BlockIdList blockIdList1 = BlockIdList.newBuilder().addAllBlockId(ImmutableList.of(1L, 2L)).build();
    LocationBlockIdListEntry listEntry1 = LocationBlockIdListEntry.newBuilder().setKey(locationProto).setValue(blockIdList1).build();
    // Prepare a lease
    GetRegisterLeasePRequest leaseRequest = GetRegisterLeasePRequest.newBuilder().setWorkerId(workerId).setBlockCount(blockIdList1.getBlockIdCount()).build();
    Optional<RegisterLease> lease = mBlockMaster.tryAcquireRegisterLease(leaseRequest);
    assertTrue(lease.isPresent());
    // Sleep for a while so that the lease expires
    SleepUtils.sleepMs(5000);
    // The lease is recycled and taken away
    GetRegisterLeasePRequest newLeaseRequest = GetRegisterLeasePRequest.newBuilder().setWorkerId(workerId + 1).setBlockCount(blockIdList1.getBlockIdCount()).build();
    Optional<RegisterLease> newLease = mBlockMaster.tryAcquireRegisterLease(newLeaseRequest);
    assertTrue(newLease.isPresent());
    RegisterWorkerPRequest request = RegisterWorkerPRequest.newBuilder().setWorkerId(workerId).addStorageTiers("MEM").putTotalBytesOnTiers("MEM", 1000L).putUsedBytesOnTiers("MEM", 0L).setOptions(RegisterWorkerPOptions.getDefaultInstance()).addCurrentBlocks(listEntry1).build();
    // Noop response observer
    Queue<Throwable> errors = new ConcurrentLinkedDeque<>();
    StreamObserver<RegisterWorkerPResponse> noopResponseObserver = new StreamObserver<RegisterWorkerPResponse>() {

        @Override
        public void onNext(RegisterWorkerPResponse response) {
        }

        @Override
        public void onError(Throwable t) {
            errors.offer(t);
        }

        @Override
        public void onCompleted() {
        }
    };
    mHandler.registerWorker(request, noopResponseObserver);
    assertEquals(1, errors.size());
    Throwable t = errors.poll();
    Assert.assertThat(t.getMessage(), containsString("does not have a lease or the lease has expired."));
    mBlockMaster.releaseRegisterLease(workerId + 1);
}

Example 40 with StreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver in project alluxio by Alluxio.

the class BlockMasterWorkerServiceHandlerTest method registerWithNoLeaseIsRejected.

@Test
public void registerWithNoLeaseIsRejected() {
    long workerId = mBlockMaster.getWorkerId(NET_ADDRESS_1);
    // Prepare LocationBlockIdListEntry objects
    BlockStoreLocation loc = new BlockStoreLocation("MEM", 0);
    BlockStoreLocationProto locationProto = BlockStoreLocationProto.newBuilder().setTierAlias(loc.tierAlias()).setMediumType(loc.mediumType()).build();
    BlockIdList blockIdList1 = BlockIdList.newBuilder().addAllBlockId(ImmutableList.of(1L, 2L)).build();
    LocationBlockIdListEntry listEntry1 = LocationBlockIdListEntry.newBuilder().setKey(locationProto).setValue(blockIdList1).build();
    RegisterWorkerPRequest request = RegisterWorkerPRequest.newBuilder().setWorkerId(workerId).addStorageTiers("MEM").putTotalBytesOnTiers("MEM", 1000L).putUsedBytesOnTiers("MEM", 0L).setOptions(RegisterWorkerPOptions.getDefaultInstance()).addCurrentBlocks(listEntry1).build();
    Queue<Throwable> errors = new ConcurrentLinkedDeque<>();
    StreamObserver<RegisterWorkerPResponse> noopResponseObserver = new StreamObserver<RegisterWorkerPResponse>() {

        @Override
        public void onNext(RegisterWorkerPResponse response) {
        }

        @Override
        public void onError(Throwable t) {
            errors.offer(t);
        }

        @Override
        public void onCompleted() {
        }
    };
    // The responseObserver should see an error
    mHandler.registerWorker(request, noopResponseObserver);
    assertEquals(1, errors.size());
    Throwable t = errors.poll();
    Assert.assertThat(t.getMessage(), containsString("does not have a lease or the lease has expired."));
}