Examples with CompletionCallback com.alibaba.graphscope.groot.CompletionCallback used on opensource projects

Example 1 with CompletionCallback

use of com.alibaba.graphscope.groot.CompletionCallback in project GraphScope by alibaba.

the class GraphWriter method writeBatch.

public void writeBatch(String requestId, String writeSession, List<WriteRequest> writeRequests, CompletionCallback<Long> callback) {
    this.pendingWriteCount.incrementAndGet();
    GraphSchema schema = snapshotCache.getSnapshotWithSchema().getGraphDef();
    OperationBatch.Builder batchBuilder = OperationBatch.newBuilder();
    for (WriteRequest writeRequest : writeRequests) {
        OperationType operationType = writeRequest.getOperationType();
        DataRecord dataRecord = writeRequest.getDataRecord();
        switch(operationType) {
            case OVERWRITE_VERTEX:
                addOverwriteVertexOperation(batchBuilder, schema, dataRecord);
                break;
            case UPDATE_VERTEX:
                addUpdateVertexOperation(batchBuilder, schema, dataRecord);
                break;
            case DELETE_VERTEX:
                addDeleteVertexOperation(batchBuilder, schema, dataRecord);
                break;
            case OVERWRITE_EDGE:
                addOverwriteEdgeOperation(batchBuilder, schema, dataRecord);
                break;
            case UPDATE_EDGE:
                addUpdateEdgeOperation(batchBuilder, schema, dataRecord);
                break;
            case DELETE_EDGE:
                addDeleteEdgeOperation(batchBuilder, schema, dataRecord);
                break;
            default:
                throw new IllegalArgumentException("Invalid operationType [" + operationType + "]");
        }
    }
    OperationBatch operationBatch = batchBuilder.build();
    int writeQueueId = getWriteQueueId(writeSession);
    int ingestorId = this.metaService.getIngestorIdForQueue(writeQueueId);
    long startTimeNano = System.nanoTime();
    this.ingestWriteClients.getClient(ingestorId).writeIngestorAsync(requestId, writeQueueId, operationBatch, new CompletionCallback<Long>() {

        @Override
        public void onCompleted(Long res) {
            long writeSnapshotId = res;
            lastWrittenSnapshotId.updateAndGet(x -> x < writeSnapshotId ? writeSnapshotId : x);
            writeRequestsTotal.addAndGet(writeRequests.size());
            finish();
            callback.onCompleted(res);
        }

        @Override
        public void onError(Throwable t) {
            finish();
            callback.onError(t);
        }

        void finish() {
            long ingestorCompleteTimeNano = System.nanoTime();
            ingestorBlockTimeNano.addAndGet(ingestorCompleteTimeNano - startTimeNano);
            pendingWriteCount.decrementAndGet();
        }
    });
}

Example 2 with CompletionCallback

use of com.alibaba.graphscope.groot.CompletionCallback in project GraphScope by alibaba.

the class IngestorRpcTest method testStoreWriteClient.

@Test
void testStoreWriteClient() {
    StoreWriteGrpc.StoreWriteStub stub = mock(StoreWriteGrpc.StoreWriteStub.class);
    StoreWriteClient client = new StoreWriteClient(stub);
    CompletionCallback callback = mock(CompletionCallback.class);
    doAnswer(invocation -> {
        StreamObserver<WriteStoreResponse> observer = invocation.getArgument(1);
        observer.onNext(WriteStoreResponse.newBuilder().setSuccess(true).build());
        return null;
    }).when(stub).writeStore(any(), any());
    client.writeStore(Arrays.asList(StoreDataBatch.newBuilder().requestId("test_req").build()), callback);
    verify(callback).onCompleted(10);
}

Example 3 with CompletionCallback

use of com.alibaba.graphscope.groot.CompletionCallback in project GraphScope by alibaba.

the class CoordinatorRpcTest method testStoreBackupClient.

@Test
void testStoreBackupClient() {
    StoreBackupGrpc.StoreBackupStub stub = mock(StoreBackupGrpc.StoreBackupStub.class);
    StoreBackupClient client = new StoreBackupClient(stub);
    StoreBackupId storeBackupId = new StoreBackupId(3);
    storeBackupId.addPartitionBackupId(0, 3);
    storeBackupId.addPartitionBackupId(1, 3);
    Map<Integer, List<Integer>> readyPartitionBackupIds = new HashMap<>();
    readyPartitionBackupIds.put(6, new ArrayList<>());
    String storeRestoreRootPath = "store_restore_path";
    doAnswer(invocation -> {
        CreateStoreBackupRequest request = invocation.getArgument(0);
        assertEquals(request.getGlobalBackupId(), 3);
        StreamObserver<CreateStoreBackupResponse> observer = invocation.getArgument(1);
        observer.onNext(CreateStoreBackupResponse.newBuilder().setStoreBackupId(storeBackupId.toProto()).build());
        observer.onError(null);
        return null;
    }).when(stub).createStoreBackup(any(), any());
    doAnswer(invocation -> {
        ClearUnavailableStoreBackupsRequest request = invocation.getArgument(0);
        assertEquals(request.getPartitionToReadyBackupIdsCount(), 1);
        assertTrue(request.getPartitionToReadyBackupIdsMap().containsKey(6));
        StreamObserver<ClearUnavailableStoreBackupsResponse> observer = invocation.getArgument(1);
        observer.onNext(ClearUnavailableStoreBackupsResponse.newBuilder().build());
        observer.onError(null);
        return null;
    }).when(stub).clearUnavailableStoreBackups(any(), any());
    doAnswer(invocation -> {
        RestoreFromStoreBackupRequest request = invocation.getArgument(0);
        assertEquals(StoreBackupId.parseProto(request.getStoreBackupId()), storeBackupId);
        assertEquals(request.getRestoreRootPath(), storeRestoreRootPath);
        StreamObserver<RestoreFromStoreBackupResponse> observer = invocation.getArgument(1);
        observer.onNext(RestoreFromStoreBackupResponse.newBuilder().build());
        observer.onError(null);
        return null;
    }).when(stub).restoreFromStoreBackup(any(), any());
    doAnswer(invocation -> {
        VerifyStoreBackupRequest request = invocation.getArgument(0);
        assertEquals(StoreBackupId.parseProto(request.getStoreBackupId()), storeBackupId);
        StreamObserver<VerifyStoreBackupResponse> observer = invocation.getArgument(1);
        observer.onNext(VerifyStoreBackupResponse.newBuilder().build());
        observer.onError(null);
        return null;
    }).when(stub).verifyStoreBackup(any(), any());
    CompletionCallback creationCallback = mock(CompletionCallback.class);
    client.createStoreBackup(3, creationCallback);
    verify(creationCallback).onCompleted(storeBackupId);
    verify(creationCallback).onError(null);
    CompletionCallback voidCallback = mock(CompletionCallback.class);
    client.clearUnavailableBackups(readyPartitionBackupIds, voidCallback);
    client.restoreFromStoreBackup(storeBackupId, storeRestoreRootPath, voidCallback);
    client.verifyStoreBackup(storeBackupId, voidCallback);
    verify(voidCallback, times(3)).onCompleted(null);
    verify(voidCallback, times(3)).onError(null);
}

Example 4 with CompletionCallback

use of com.alibaba.graphscope.groot.CompletionCallback in project GraphScope by alibaba.

the class CoordinatorRpcTest method testIngestorSnapshotClient.

@Test
void testIngestorSnapshotClient() {
    IngestorSnapshotGrpc.IngestorSnapshotStub stub = mock(IngestorSnapshotGrpc.IngestorSnapshotStub.class);
    IngestorSnapshotClient client = new IngestorSnapshotClient(stub);
    CompletionCallback callback = mock(CompletionCallback.class);
    doAnswer(invocation -> {
        AdvanceIngestSnapshotIdRequest req = invocation.getArgument(0);
        assertEquals(req.getSnapshotId(), 10L);
        StreamObserver<AdvanceIngestSnapshotIdResponse> observer = invocation.getArgument(1);
        observer.onNext(AdvanceIngestSnapshotIdResponse.newBuilder().setPreviousSnapshotId(8L).build());
        observer.onError(null);
        return null;
    }).when(stub).advanceIngestSnapshotId(any(), any());
    client.advanceIngestSnapshotId(10L, callback);
    verify(callback).onCompleted(8L);
    verify(callback).onError(null);
}

Example 5 with CompletionCallback

use of com.alibaba.graphscope.groot.CompletionCallback in project GraphScope by alibaba.

the class CoordinatorRpcTest method testFrontendSnapshotClient.

@Test
void testFrontendSnapshotClient() {
    FrontendSnapshotGrpc.FrontendSnapshotStub stub = mock(FrontendSnapshotGrpc.FrontendSnapshotStub.class);
    FrontendSnapshotClient client = new FrontendSnapshotClient(stub);
    GraphDef graphDef = GraphDef.newBuilder().build();
    doAnswer(invocationOnMock -> {
        AdvanceQuerySnapshotRequest req = invocationOnMock.getArgument(0);
        StreamObserver<AdvanceQuerySnapshotResponse> observer = invocationOnMock.getArgument(1);
        assertEquals(req.getGraphDef(), graphDef.toProto());
        assertEquals(req.getSnapshotId(), 10L);
        observer.onNext(AdvanceQuerySnapshotResponse.newBuilder().setPreviousSnapshotId(9L).build());
        return null;
    }).when(stub).advanceQuerySnapshot(any(), any());
    CompletionCallback completionCallback = mock(CompletionCallback.class);
    client.advanceQuerySnapshot(10L, graphDef, completionCallback);
    verify(completionCallback).onCompleted(9L);
}