Examples with BatchElement com.palantir.atlasdb.autobatch.BatchElement used on opensource projects

Example 6 with BatchElement

use of com.palantir.atlasdb.autobatch.BatchElement in project atlasdb by palantir.

the class MultiClientTransactionStarterTest method servesRequestsAsSoonAsResponseIsReceived.

@Test
public void servesRequestsAsSoonAsResponseIsReceived() {
    Namespace namespace = Namespace.of("Test" + UUID.randomUUID());
    UUID requestorId = UUID.randomUUID();
    BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>> requestToBeServed = batchElementForNamespace(namespace, 1);
    BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>> requestNotToBeServed = batchElementForNamespace(namespace, PARTITIONED_TIMESTAMPS_LIMIT_PER_SERVER_CALL * 5);
    ImmutableList<BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>>> requests = ImmutableList.of(requestToBeServed, requestNotToBeServed);
    Map<Namespace, ConjureStartTransactionsResponse> responseMap = startTransactionsResponse(requests, requestorId);
    when(timelockService.startTransactions(any())).thenReturn(responseMap).thenThrow(EXCEPTION);
    assertThatThrownBy(() -> processBatch(timelockService, requestorId, requests)).isEqualTo(EXCEPTION);
    // assert first request is served even if server throws on next request
    assertSanityOfRequestBatch(ImmutableList.of(requestToBeServed), ImmutableMap.of(namespace, ImmutableList.of(responseMap.get(namespace))));
    assertThat(requestNotToBeServed.result().isDone()).isFalse();
    LockCleanupService relevantLockCleanupService = LOCK_CLEANUP_SERVICE_MAP.get(namespace);
    verify(relevantLockCleanupService).refreshLockLeases(any());
    verify(relevantLockCleanupService).unlock(any());
    /*
         * For one space, all requests are accumulated, and we attempt to fetch as many respones as possible from the
         * server. Concretely, the queue here contains 25 demands for a response. The first one succeeds, but the
         * next four fail (and each batch size is five), and thus we get four calls to clean up.
         */
    verify(NAMESPACE_CACHE_MAP.get(namespace), times(4)).removeTransactionStateFromCache(anyLong());
}

Example 7 with BatchElement

use of com.palantir.atlasdb.autobatch.BatchElement in project atlasdb by palantir.

the class MultiClientTransactionStarterTest method shouldNotFreeResourcesIfRequestIsServed.

@Test
public void shouldNotFreeResourcesIfRequestIsServed() {
    Namespace alpha = Namespace.of("alpha" + UUID.randomUUID());
    Namespace beta = Namespace.of("beta" + UUID.randomUUID());
    BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>> requestForAlpha = batchElementForNamespace(alpha, PARTITIONED_TIMESTAMPS_LIMIT_PER_SERVER_CALL - 1);
    BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>> requestForBeta = batchElementForNamespace(beta, PARTITIONED_TIMESTAMPS_LIMIT_PER_SERVER_CALL * 5);
    UUID requestorId = UUID.randomUUID();
    List<BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>>> requests = ImmutableList.of(requestForAlpha, requestForBeta);
    Map<Namespace, ConjureStartTransactionsResponse> responseMap = startTransactionsResponse(requests, requestorId);
    when(timelockService.startTransactions(any())).thenReturn(responseMap).thenThrow(EXCEPTION);
    assertThatThrownBy(() -> processBatch(timelockService, requestorId, requests)).isEqualTo(EXCEPTION);
    // assert requests made by client alpha are served
    assertSanityOfRequestBatch(ImmutableList.of(requestForAlpha), ImmutableMap.of(alpha, ImmutableList.of(responseMap.get(alpha))));
    verify(LOCK_CLEANUP_SERVICE_MAP.get(alpha), never()).unlock(any());
    verify(LOCK_CLEANUP_SERVICE_MAP.get(beta)).refreshLockLeases(any());
    verify(LOCK_CLEANUP_SERVICE_MAP.get(beta)).unlock(any());
    verify(NAMESPACE_CACHE_MAP.get(alpha), never()).removeTransactionStateFromCache(anyLong());
    // The size of each batch is 5 here, and thus for a single batch we need to clean up five times
    verify(NAMESPACE_CACHE_MAP.get(beta), times(5)).removeTransactionStateFromCache(anyLong());
}

Example 8 with BatchElement

use of com.palantir.atlasdb.autobatch.BatchElement in project atlasdb by palantir.

the class MultiClientTransactionStarter method getResponseHandlers.

private static Map<Namespace, ResponseHandler> getResponseHandlers(List<BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>>> batch) {
    Map<Namespace, ResponseHandler> responseHandlers = new HashMap<>();
    for (BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>> element : batch) {
        NamespaceAndRequestParams requestParams = element.argument();
        Namespace namespace = requestParams.namespace();
        responseHandlers.computeIfAbsent(namespace, _unused -> new ResponseHandler(requestParams.params().lockCleanupService(), requestParams.params().cache())).addPendingFuture(SettableResponse.of(requestParams.params().numTransactions(), element.result()));
    }
    return responseHandlers;
}

Example 9 with BatchElement

use of com.palantir.atlasdb.autobatch.BatchElement in project atlasdb by palantir.

the class MultiClientTransactionStarterTest method shouldNotFreeResourcesWithinNamespaceIfRequestIsServed.

@Test
public void shouldNotFreeResourcesWithinNamespaceIfRequestIsServed() {
    Namespace omega = Namespace.of("omega" + UUID.randomUUID());
    BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>> requestForOmega = batchElementForNamespace(omega, PARTITIONED_TIMESTAMPS_LIMIT_PER_SERVER_CALL - 1);
    BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>> secondRequestForOmega = batchElementForNamespace(omega, 2);
    UUID requestorId = UUID.randomUUID();
    List<BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>>> requests = ImmutableList.of(requestForOmega, secondRequestForOmega);
    Map<Namespace, ConjureStartTransactionsResponse> responseMap = startTransactionsResponse(requests, requestorId);
    when(timelockService.startTransactions(any())).thenReturn(responseMap).thenThrow(EXCEPTION);
    assertThatThrownBy(() -> processBatch(timelockService, requestorId, requests)).isEqualTo(EXCEPTION);
    // assert the first request made by client omega is served
    assertSanityOfRequestBatch(ImmutableList.of(requestForOmega), ImmutableMap.of(omega, ImmutableList.of(responseMap.get(omega))));
    @SuppressWarnings({ "unchecked", "rawtypes" }) ArgumentCaptor<Set<LockToken>> refreshArgumentCaptor = (ArgumentCaptor<Set<LockToken>>) ArgumentCaptor.forClass((Class) Set.class);
    verify(LOCK_CLEANUP_SERVICE_MAP.get(omega)).refreshLockLeases(refreshArgumentCaptor.capture());
    verify(LOCK_CLEANUP_SERVICE_MAP.get(omega)).unlock(eq(Collections.emptySet()));
    verify(NAMESPACE_CACHE_MAP.get(omega)).removeTransactionStateFromCache(anyLong());
    Set<LockToken> refreshedTokens = refreshArgumentCaptor.getValue();
    LockToken tokenShare = Futures.getUnchecked(requestForOmega.result()).get(0).immutableTimestamp().getLock();
    assertThat(tokenShare).isInstanceOf(LockTokenShare.class).satisfies(token -> {
        LockTokenShare share = ((LockTokenShare) token);
        assertThat(share.sharedLockToken()).isIn(refreshedTokens);
    });
}

Example 10 with BatchElement

use of com.palantir.atlasdb.autobatch.BatchElement in project atlasdb by palantir.

the class MultiClientTransactionStarterTest method assertSanityOfRequestBatch.

private void assertSanityOfRequestBatch(List<BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>>> requestsForClients, Map<Namespace, List<ConjureStartTransactionsResponse>> responseMap) {
    assertCompletedWithCorrectNumberOfTransactions(requestsForClients);
    Map<Namespace, List<BatchElement<NamespaceAndRequestParams, List<StartIdentifiedAtlasDbTransactionResponse>>>> partitionedResponses = requestsForClients.stream().collect(Collectors.groupingBy(e -> e.argument().namespace()));
    responseMap.forEach((namespace, responses) -> {
        int startInd = 0;
        List<StartIdentifiedAtlasDbTransactionResponse> startedTransactions = partitionedResponses.get(namespace).stream().map(response -> Futures.getUnchecked(response.result())).flatMap(Collection::stream).collect(Collectors.toList());
        for (ConjureStartTransactionsResponse conjureResponse : responses) {
            int toIndex = Math.min(startInd + conjureResponse.getTimestamps().count(), startedTransactions.size());
            List<StartIdentifiedAtlasDbTransactionResponse> responseList = startedTransactions.subList(startInd, toIndex);
            startInd = toIndex;
            assertThat(responseList).satisfies(StartTransactionsTestUtils::assertThatStartTransactionResponsesAreUnique).allSatisfy(startTxnResponse -> StartTransactionsTestUtils.assertDerivableFromBatchedResponse(startTxnResponse, conjureResponse));
        }
    });
}