Examples with StoreWriter com.alibaba.graphscope.groot.ingestor.StoreWriter used on opensource projects

Example 1 with StoreWriter

use of com.alibaba.graphscope.groot.ingestor.StoreWriter in project GraphScope by alibaba.

the class BatchSenderTest method testSend.

@Test
void testSend() throws InterruptedException {
    Configs configs = Configs.newBuilder().put(CommonConfig.STORE_NODE_COUNT.getKey(), "1").build();
    MetaService mockMetaService = mock(MetaService.class);
    when(mockMetaService.getPartitionCount()).thenReturn(2);
    when(mockMetaService.getStoreIdByPartition(anyInt())).thenReturn(0);
    StoreWriter mockStoreWriter = mock(StoreWriter.class);
    String requestId = "test_batch_sender";
    int queueId = 0;
    long snapshotId = 10L;
    long offset = 50L;
    LabelId labelId = new LabelId(0);
    OperationBlob writeVertexBlob1 = new OverwriteVertexOperation(new VertexId(0L), labelId, Collections.EMPTY_MAP).toBlob();
    OperationBlob writeVertexBlob2 = new OverwriteVertexOperation(new VertexId(1L), labelId, Collections.EMPTY_MAP).toBlob();
    OperationBatch batch = OperationBatch.newBuilder().addOperationBlob(writeVertexBlob1).addOperationBlob(writeVertexBlob2).addOperationBlob(OperationBlob.MARKER_OPERATION_BLOB).build();
    CountDownLatch latch = new CountDownLatch(1);
    doAnswer(invocationOnMock -> {
        StoreDataBatch storeBatch = invocationOnMock.getArgument(1);
        CompletionCallback callback = invocationOnMock.getArgument(2);
        assertAll(() -> assertEquals(storeBatch.getRequestId(), requestId), () -> assertEquals(storeBatch.getQueueId(), queueId), () -> assertEquals(storeBatch.getSnapshotId(), snapshotId), () -> assertEquals(storeBatch.getOffset(), offset), () -> assertEquals(storeBatch.getDataBatch().size(), 2));
        List<Map<Integer, OperationBatch>> dataBatch = storeBatch.getDataBatch();
        Map<Integer, OperationBatch> partitionToBatch = dataBatch.get(0);
        assertAll(() -> assertEquals(partitionToBatch.get(0).getOperationBlob(0), writeVertexBlob1), () -> assertEquals(partitionToBatch.get(1).getOperationBlob(0), writeVertexBlob2));
        assertEquals(dataBatch.get(1).get(-1).getOperationBlob(0), OperationBlob.MARKER_OPERATION_BLOB);
        callback.onCompleted(0);
        latch.countDown();
        return null;
    }).when(mockStoreWriter).write(anyInt(), any(), any());
    BatchSender batchSender = new BatchSender(configs, mockMetaService, mockStoreWriter, new MetricsCollector(configs));
    batchSender.start();
    batchSender.asyncSendWithRetry(requestId, queueId, snapshotId, offset, batch);
    assertTrue(latch.await(5L, TimeUnit.SECONDS));
    batchSender.stop();
}

Example 2 with StoreWriter

use of com.alibaba.graphscope.groot.ingestor.StoreWriter in project GraphScope by alibaba.

the class IngestServiceTest method testIngestService.

@Test
void testIngestService() {
    Configs configs = Configs.newBuilder().put(CommonConfig.NODE_IDX.getKey(), "0").put(CommonConfig.STORE_NODE_COUNT.getKey(), "1").put(IngestorConfig.INGESTOR_CHECK_PROCESSOR_INTERVAL_MS.getKey(), "100").build();
    MockDiscovery mockDiscovery = new MockDiscovery();
    MetaService mockMetaService = mock(MetaService.class);
    when(mockMetaService.getQueueIdsForIngestor(0)).thenReturn(Arrays.asList(0));
    LogService mockLogService = mock(LogService.class);
    IngestProgressFetcher mockIngestProgressFetcher = mock(IngestProgressFetcher.class);
    when(mockIngestProgressFetcher.getTailOffsets(Arrays.asList(0))).thenReturn(Arrays.asList(50L));
    StoreWriter mockStoreWriter = mock(StoreWriter.class);
    IngestService spyIngestService = spy(new IngestService(configs, mockDiscovery, mockMetaService, mockLogService, mockIngestProgressFetcher, mockStoreWriter, new MetricsCollector(configs)));
    IngestProcessor mockIngestProcessor = mock(IngestProcessor.class);
    doReturn(mockIngestProcessor).when(spyIngestService).makeIngestProcessor(any(), any(), any(), eq(0), any(), any());
    spyIngestService.start();
    verify(mockIngestProcessor, never()).start();
    mockDiscovery.addNode(RoleType.STORE, Collections.singletonMap(0, null));
    verify(mockIngestProcessor, timeout(5000L)).setTailOffset(50L);
    verify(mockIngestProcessor, timeout(5000L)).start();
    spyIngestService.advanceIngestSnapshotId(5L, null);
    verify(mockIngestProcessor).ingestBatch(eq("marker"), eq(IngestService.MARKER_BATCH), any());
    mockDiscovery.removeNode(RoleType.STORE, Collections.singletonMap(0, null));
    verify(mockIngestProcessor, timeout(5000L).times(2)).stop();
    spyIngestService.stop();
}