Examples with BlockWorkerClient alluxio.client.block.stream.BlockWorkerClient used on opensource projects

Example 6 with BlockWorkerClient

use of alluxio.client.block.stream.BlockWorkerClient in project alluxio by Alluxio.

the class AlluxioFileInStreamTest method before.

/**
 * Sets up the context and streams before a test runs.
 */
@Before
public void before() throws Exception {
    mInfo = new FileInfo().setBlockSizeBytes(BLOCK_LENGTH).setLength(mFileSize);
    ClientTestUtils.setSmallBufferSizes(mConf);
    mConf.set(PropertyKey.USER_BLOCK_READ_RETRY_SLEEP_MIN, "1ms");
    mConf.set(PropertyKey.USER_BLOCK_READ_RETRY_SLEEP_MAX, "5ms");
    mConf.set(PropertyKey.USER_BLOCK_READ_RETRY_MAX_DURATION, "1s");
    BlockWorkerClient client = mock(BlockWorkerClient.class);
    doNothing().when(client).cache(any());
    mContext = mock(FileSystemContext.class);
    when(mContext.getClientContext()).thenReturn(ClientContext.create(mConf));
    when(mContext.getClusterConf()).thenReturn(mConf);
    when(mContext.getPathConf(any(AlluxioURI.class))).thenReturn(mConf);
    when(mContext.getNodeLocalWorker()).thenReturn(new WorkerNetAddress());
    when(mContext.getCachedWorkers()).thenReturn(new ArrayList<>());
    when(mContext.acquireBlockWorkerClient(any())).thenReturn(new CloseableResource<BlockWorkerClient>(client) {

        @Override
        public void closeResource() {
        }
    });
    mBlockStore = mock(AlluxioBlockStore.class);
    PowerMockito.mockStatic(AlluxioBlockStore.class);
    PowerMockito.when(AlluxioBlockStore.create(mContext)).thenReturn(mBlockStore);
    // Set up BufferedBlockInStreams and caching streams
    mInStreams = new ArrayList<>();
    List<Long> blockIds = new ArrayList<>();
    List<FileBlockInfo> fileBlockInfos = new ArrayList<>();
    for (int i = 0; i < mNumBlocks; i++) {
        blockIds.add((long) i);
        FileBlockInfo fbInfo = new FileBlockInfo().setBlockInfo(new BlockInfo().setBlockId(i));
        fileBlockInfos.add(fbInfo);
        final byte[] input = BufferUtils.getIncreasingByteArray((int) (i * BLOCK_LENGTH), (int) getBlockLength(i));
        mInStreams.add(new TestBlockInStream(input, i, input.length, false, mBlockSource));
        when(mContext.getCachedWorkers()).thenReturn(Arrays.asList(new BlockWorkerInfo(new WorkerNetAddress(), 0, 0)));
        when(mBlockStore.getInStream(eq((long) i), any(InStreamOptions.class), any())).thenAnswer(invocation -> {
            long blockId = (Long) invocation.getArguments()[0];
            return mInStreams.get((int) blockId).isClosed() ? new TestBlockInStream(input, blockId, input.length, false, mBlockSource) : mInStreams.get((int) blockId);
        });
        when(mBlockStore.getInStream(eq(new BlockInfo().setBlockId(i)), any(InStreamOptions.class), any())).thenAnswer(invocation -> {
            long blockId = ((BlockInfo) invocation.getArguments()[0]).getBlockId();
            return mInStreams.get((int) blockId).isClosed() ? new TestBlockInStream(input, blockId, input.length, false, mBlockSource) : mInStreams.get((int) blockId);
        });
    }
    mInfo.setBlockIds(blockIds);
    mInfo.setFileBlockInfos(fileBlockInfos).setReplicationMax(1);
    mStatus = new URIStatus(mInfo);
    OpenFilePOptions readOptions = OpenFilePOptions.newBuilder().setReadType(ReadPType.CACHE_PROMOTE).build();
    mTestStream = new AlluxioFileInStream(mStatus, new InStreamOptions(mStatus, readOptions, mConf), mContext);
}

Example 7 with BlockWorkerClient

use of alluxio.client.block.stream.BlockWorkerClient in project alluxio by Alluxio.

the class JobUtils method loadThroughCacheRequest.

private static void loadThroughCacheRequest(URIStatus status, FileSystemContext context, long blockId, AlluxioConfiguration conf, WorkerNetAddress localNetAddress) throws IOException {
    AlluxioBlockStore blockStore = AlluxioBlockStore.create(context);
    OpenFilePOptions openOptions = OpenFilePOptions.newBuilder().setReadType(ReadPType.CACHE).build();
    InStreamOptions inOptions = new InStreamOptions(status, openOptions, conf);
    BlockLocationPolicy policy = BlockLocationPolicy.Factory.create(LocalFirstPolicy.class.getCanonicalName(), conf);
    inOptions.setUfsReadLocationPolicy(policy);
    Protocol.OpenUfsBlockOptions openUfsBlockOptions = inOptions.getOpenUfsBlockOptions(blockId);
    BlockInfo info = Preconditions.checkNotNull(status.getBlockInfo(blockId));
    long blockLength = info.getLength();
    Pair<WorkerNetAddress, BlockInStream.BlockInStreamSource> dataSourceAndType = blockStore.getDataSourceAndType(status.getBlockInfo(blockId), status, policy, ImmutableMap.of());
    WorkerNetAddress dataSource = dataSourceAndType.getFirst();
    String host = dataSource.getHost();
    // to establish the connection.
    if (!dataSource.getContainerHost().equals("")) {
        host = dataSource.getContainerHost();
    }
    CacheRequest request = CacheRequest.newBuilder().setBlockId(blockId).setLength(blockLength).setOpenUfsBlockOptions(openUfsBlockOptions).setSourceHost(host).setSourcePort(dataSource.getDataPort()).build();
    try (CloseableResource<BlockWorkerClient> blockWorker = context.acquireBlockWorkerClient(localNetAddress)) {
        blockWorker.get().cache(request);
    } catch (Exception e) {
        throw new IOException(e);
    }
}