@Override public void onSuccess(Response<PermitAllocation> result) { BatchedPermitsRequester.this.retries = 0; BatchedPermitsRequester.this.callbackCounter.incrementAndGet(); BatchedPermitsRequester.this.lock.lock(); try { PermitAllocation allocation = result.getEntity(); log.debug("Received permit allocation " + allocation); Long retryDelay = allocation.getMinRetryDelayMillis(GetMode.NULL); if (retryDelay != null) { BatchedPermitsRequester.this.retryStatus.blockRetries(retryDelay, null); } if (allocation.getPermits() > 0) { BatchedPermitsRequester.this.permitBatchContainer.addPermitAllocation(allocation); } BatchedPermitsRequester.this.requestSemaphore.release(); if (allocation.getPermits() > 0) { BatchedPermitsRequester.this.newPermitsAvailable.signalAll(); } } finally { try { this.timerContext.close(); } catch (IOException ioe) { // Do nothing } BatchedPermitsRequester.this.lock.unlock(); } }
@Override public void onSuccess(Response<RT> responseToGet) { batch.entries().stream().forEach(entry -> { Request request = entry.getKey().getRequest(); if (request instanceof GetRequest) { entry.getValue().getPromise().done(new ResponseImpl<>(responseToGet, responseToGet.getEntity())); } else { entry.getValue().getPromise().fail(unsupportedGetRequestType(request)); } }); }
private void successBatchGet(BatchGetRequest request, Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch, Entry<RestRequestBatchKey, BatchEntry<Response<Object>>> entry, final ProtocolVersion version) { Set<String> ids = (Set<String>) request.getObjectIds().stream() .map(o -> BatchResponse.keyToString(o, version)) .collect(Collectors.toSet()); DataMap dm = filterIdsInBatchResult(responseToBatch.getEntity().data(), ids); BatchResponse br = new BatchResponse<>(dm, request.getResponseDecoder().getEntityClass()); Response rsp = new ResponseImpl(responseToBatch, br); entry.getValue().getPromise().done(rsp); }
/** * Get the latest flow status * @param flowId identifier of flow status to get * @return a {@link FlowStatus} with the flow status * @throws RemoteInvocationException */ public FlowStatus getLatestFlowStatus(FlowId flowId) throws RemoteInvocationException { LOG.debug("getFlowConfig with groupName " + flowId.getFlowGroup() + " flowName " + flowId.getFlowName()); FindRequest<FlowStatus> findRequest = _flowstatusesRequestBuilders.findByLatestFlowStatus().flowIdParam(flowId).build(); Response<CollectionResponse<FlowStatus>> response = _restClient.get().sendRequest(findRequest).getResponse(); List<FlowStatus> flowStatusList = response.getEntity().getElements(); if (flowStatusList.isEmpty()) { return null; } else { Preconditions.checkArgument(flowStatusList.size() == 1); return flowStatusList.get(0); } }
private static <K, RT extends RecordTemplate> Response<RT> unbatchResponse(BatchGetEntityRequest<K, RT> request, Response<BatchKVResponse<K, EntityResponse<RT>>> batchResponse, Object id) throws RemoteInvocationException { final BatchKVResponse<K, EntityResponse<RT>> batchEntity = batchResponse.getEntity(); final ErrorResponse errorResponse = batchEntity.getErrors().get(id); if (errorResponse != null) { throw new RestLiResponseException(errorResponse); } final EntityResponse<RT> entityResponse = batchEntity.getResults().get(id); if (entityResponse != null) { final RT entityResult = entityResponse.getEntity(); if (entityResult != null) { return new ResponseImpl<>(batchResponse, entityResult); } } LOGGER.debug("No result or error for base URI : {}, id: {}. Verify that the batchGet endpoint returns response keys that match batchGet request IDs.", request.getBaseUriTemplate(), id); throw NOT_FOUND_EXCEPTION; }
protected Task<String> toMessage(Task<Response<Greeting>> greeting) { return greeting.map("toMessage", g -> g.getEntity().getMessage()); }
@Override public void sendRequest(PermitRequest request, Callback<Response<PermitAllocation>> callback) { this.requestList.add(request); PermitAllocation permitAllocation = new PermitAllocation(); permitAllocation.setPermits(request.getPermits()); permitAllocation.setExpiration(Long.MAX_VALUE); Response<PermitAllocation> response = Mockito.mock(Response.class); Mockito.when(response.getEntity()).thenReturn(permitAllocation); callback.onSuccess(response); } }
@SuppressWarnings({ "deprecation" }) private void successBatchGetEntity(BatchGetEntityRequest request, Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch, Entry<RestRequestBatchKey, BatchEntry<Response<Object>>> entry, final ProtocolVersion version) { Set<String> ids = (Set<String>) request.getObjectIds().stream() .map(o -> BatchResponse.keyToString(o, version)) .collect(Collectors.toSet()); DataMap dm = filterIdsInBatchResult(responseToBatch.getEntity().data(), ids); BatchKVResponse br = new BatchEntityResponse<>(dm, request.getResourceSpec().getKeyType(), request.getResourceSpec().getValueType(), request.getResourceSpec().getKeyParts(), request.getResourceSpec().getComplexKeyType(), version); Response rsp = new ResponseImpl(responseToBatch, br); entry.getValue().getPromise().done(rsp); }
@SuppressWarnings({ "deprecation" }) private void successBatchGetKV(BatchGetKVRequest request, Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch, Entry<RestRequestBatchKey, BatchEntry<Response<Object>>> entry, final ProtocolVersion version) { Set<String> ids = (Set<String>) request.getObjectIds().stream() .map(o -> BatchResponse.keyToString(o, version)) .collect(Collectors.toSet()); DataMap dm = filterIdsInBatchResult(responseToBatch.getEntity().data(), ids); BatchKVResponse br = new BatchKVResponse(dm, request.getResourceSpec().getKeyType(), request.getResourceSpec().getValueType(), request.getResourceSpec().getKeyParts(), request.getResourceSpec().getComplexKeyType(), version); Response rsp = new ResponseImpl(responseToBatch, br); entry.getValue().getPromise().done(rsp); }
/** * Get a flow configuration * @param flowId identifier of flow configuration to get * @return a {@link FlowConfig} with the flow configuration * @throws RemoteInvocationException */ public FlowConfig getFlowConfig(FlowId flowId) throws RemoteInvocationException { LOG.debug("getFlowConfig with groupName " + flowId.getFlowGroup() + " flowName " + flowId.getFlowName()); GetRequest<FlowConfig> getRequest = _flowconfigsRequestBuilders.get() .id(new ComplexResourceKey<>(flowId, new EmptyRecord())).build(); Response<FlowConfig> response = _restClient.get().sendRequest(getRequest).getResponse(); return response.getEntity(); }
/** * Get a flow status * @param flowStatusId identifier of flow status to get * @return a {@link FlowStatus} with the flow status * @throws RemoteInvocationException */ public FlowStatus getFlowStatus(FlowStatusId flowStatusId) throws RemoteInvocationException { LOG.debug("getFlowConfig with groupName " + flowStatusId.getFlowGroup() + " flowName " + flowStatusId.getFlowName()); GetRequest<FlowStatus> getRequest = _flowstatusesRequestBuilders.get() .id(new ComplexResourceKey<>(flowStatusId, new EmptyRecord())).build(); Response<FlowStatus> response = _restClient.get().sendRequest(getRequest).getResponse(); return response.getEntity(); }
/** * Get a flow configuration * @param flowId identifier of flow configuration to get * @return a {@link FlowConfig} with the flow configuration * @throws RemoteInvocationException */ public FlowConfig getFlowConfig(FlowId flowId) throws RemoteInvocationException { LOG.debug("getFlowConfig with groupName " + flowId.getFlowGroup() + " flowName " + flowId.getFlowName()); GetRequest<FlowConfig> getRequest = _flowconfigsV2RequestBuilders.get() .id(new ComplexResourceKey<>(flowId, new FlowStatusId())).build(); Response<FlowConfig> response = _restClient.get().sendRequest(getRequest).getResponse(); return response.getEntity(); }
@Builder(builderMethodName = "satisfyRequestBuilder", buildMethodName = "satisfy") public static void satisfyRequest(RequestAndCallback requestAndCallback, long expiration) { PermitAllocation allocation = new PermitAllocation(); allocation.setPermits(requestAndCallback.getRequest().getPermits()); allocation.setExpiration(expiration > 0 ? expiration : Long.MAX_VALUE); Response<PermitAllocation> response = Mockito.mock(Response.class); Mockito.when(response.getEntity()).thenReturn(allocation); requestAndCallback.getCallback().onSuccess(response); }
@Test public void testGetSubResourceRequests() { Tuple2Task<Response<Message>,Response<Message>> task = Task.par(associationsGet("a", "b", "x"), associationsGet("a", "b", "y")); if (expectBatching()) { runAndWaitException(task, RestLiResponseException.class); assertTrue(((RestLiResponseException)task.getError()).getServiceErrorMessage().contains("associationsSub?ids=List(x,y)")); } else { runAndWait(getTestClassName() + ".testGetSubResourceRequests", task); assertEquals(task.get()._1().getEntity().getMessage(), "b"); assertEquals(task.get()._1().getEntity().getId(), "a"); assertEquals(task.get()._2().getEntity().getMessage(), "b"); assertEquals(task.get()._2().getEntity().getId(), "a"); } }
@Test public void testGetSubResourceRequestsOverrides() { Tuple2Task<Response<Message>,Response<Message>> task = Task.par(associationsGet("a", "b", "x", overrides()), associationsGet("a", "b", "y", overrides())); if (expectBatchingOverrides()) { runAndWaitException(task, RestLiResponseException.class); assertTrue(((RestLiResponseException)task.getError()).getServiceErrorMessage().contains("associationsSub?ids=List(x,y)")); } else { runAndWait(getTestClassName() + ".testGetSubResourceRequestsOverrides", task); assertEquals(task.get()._1().getEntity().getMessage(), "b"); assertEquals(task.get()._1().getEntity().getId(), "a"); assertEquals(task.get()._2().getEntity().getMessage(), "b"); assertEquals(task.get()._2().getEntity().getId(), "a"); } }
private Object remove404(Object o) { if (o instanceof Response) { Response r = (Response) o; Object entity = r.getEntity(); if (entity instanceof BatchEntityResponse) { BatchEntityResponse ber = (BatchEntityResponse) entity; DataMap data = ber.data(); DataMap errors = (DataMap) data.getDataMap("errors"); Set<String> keys = new HashSet<>(errors.keySet()); keys.forEach(key -> { DataMap error = errors.getDataMap(key); if (error.getInteger("status").equals(404)) { errors.remove(key); } }); } } return o; }
@Override public void run() { try { while (true) { RequestAndCallback requestAndCallback = MockRequester.this.requestAndCallbackQueue.take(); long nanoTime = System.nanoTime(); long delayNanos = requestAndCallback.getProcessAfterNanos() - nanoTime; if (delayNanos > 0) { Thread.sleep(TimeUnit.NANOSECONDS.toMillis(delayNanos)); } try { PermitAllocation allocation = MockRequester.this.limiterServer.getSync(new ComplexResourceKey<>(requestAndCallback.getRequest(), new EmptyRecord())); Response<PermitAllocation> response = Mockito.mock(Response.class); Mockito.when(response.getEntity()).thenReturn(allocation); requestAndCallback.getCallback().onSuccess(response); } catch (RestLiServiceException rexc) { RestLiResponseException returnException = Mockito.mock(RestLiResponseException.class); Mockito.when(returnException.getStatus()).thenReturn(rexc.getStatus().getCode()); requestAndCallback.getCallback().onError(returnException); } } } catch (Throwable t) { log.error("Error", t); throw new RuntimeException(t); } } }
/** * Get a {@link org.apache.gobblin.rest.JobExecutionQueryResult} for a {@link org.apache.gobblin.rest.JobExecutionQuery}. * * @param query a {@link org.apache.gobblin.rest.JobExecutionQuery} * @return a {@link org.apache.gobblin.rest.JobExecutionQueryResult} * @throws RemoteInvocationException */ public JobExecutionQueryResult get(JobExecutionQuery query) throws RemoteInvocationException { GetRequest<JobExecutionQueryResult> getRequest = new JobExecutionsBuilders().get() .id(new ComplexResourceKey<JobExecutionQuery, EmptyRecord>(query, new EmptyRecord())).build(); Response<JobExecutionQueryResult> response = this.restClient.sendRequest(getRequest, ErrorHandlingBehavior.TREAT_SERVER_ERROR_AS_SUCCESS).getResponse(); return response.getEntity(); }
private PermitAllocation getPermitAllocation(PermitRequest permitRequest, RestClient restClient, PermitsGetRequestBuilder getBuilder) throws Exception { Request<PermitAllocation> request = getBuilder.id(new ComplexResourceKey<>(permitRequest, new EmptyRecord())).build(); ResponseFuture<PermitAllocation> responseFuture = restClient.sendRequest(request); Response<PermitAllocation> response = responseFuture.getResponse(); return response.getEntity(); }
@Override public T getResponseEntity() throws RemoteInvocationException { return getResponse().getEntity(); }