Examples with GetDataRequest org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetDataRequest used on opensource projects

Example 1 with GetDataRequest

use of org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetDataRequest in project beam by apache.

the class FakeWindmillServer method getData.

@Override
public Windmill.GetDataResponse getData(Windmill.GetDataRequest request) {
    LOG.info("getDataRequest: {}", request.toString());
    validateGetDataRequest(request);
    ++numGetDataRequests;
    GetDataResponse response;
    Function<GetDataRequest, GetDataResponse> responseFn = dataToOffer.poll();
    if (responseFn == null) {
        response = Windmill.GetDataResponse.newBuilder().build();
    } else {
        response = responseFn.apply(request);
        try {
            // Sleep for a little bit to ensure that *-windmill-read state-sampled counters
            // show up.
            sleepMillis(500);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }
    LOG.debug("getDataResponse: {}", response.toString());
    return response;
}

Example 2 with GetDataRequest

use of org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetDataRequest in project beam by apache.

the class FakeWindmillServer method getDataStream.

@Override
public GetDataStream getDataStream() {
    Instant startTime = Instant.now();
    return new GetDataStream() {

        @Override
        public Windmill.KeyedGetDataResponse requestKeyedData(String computation, KeyedGetDataRequest request) {
            Windmill.GetDataRequest getDataRequest = GetDataRequest.newBuilder().addRequests(ComputationGetDataRequest.newBuilder().setComputationId(computation).addRequests(request).build()).build();
            GetDataResponse getDataResponse = getData(getDataRequest);
            if (getDataResponse.getDataList().isEmpty()) {
                return null;
            }
            assertEquals(1, getDataResponse.getDataCount());
            if (getDataResponse.getData(0).getDataList().isEmpty()) {
                return null;
            }
            assertEquals(1, getDataResponse.getData(0).getDataCount());
            return getDataResponse.getData(0).getData(0);
        }

        @Override
        public Windmill.GlobalData requestGlobalData(Windmill.GlobalDataRequest request) {
            Windmill.GetDataRequest getDataRequest = GetDataRequest.newBuilder().addGlobalDataFetchRequests(request).build();
            GetDataResponse getDataResponse = getData(getDataRequest);
            if (getDataResponse.getGlobalDataList().isEmpty()) {
                return null;
            }
            assertEquals(1, getDataResponse.getGlobalDataCount());
            return getDataResponse.getGlobalData(0);
        }

        @Override
        public void refreshActiveWork(Map<String, List<KeyedGetDataRequest>> active) {
        }

        @Override
        public void close() {
        }

        @Override
        public boolean awaitTermination(int time, TimeUnit unit) {
            return true;
        }

        @Override
        public Instant startTime() {
            return startTime;
        }
    };
}