Examples with StageControl org.neo4j.unsafe.impl.batchimport.staging.StageControl used on opensource projects

Example 1 with StageControl

use of org.neo4j.unsafe.impl.batchimport.staging.StageControl in project neo4j by neo4j.

the class EncodeGroupsStepTest method shouldEncodeGroupChains.

@SuppressWarnings("unchecked")
@Test
public void shouldEncodeGroupChains() throws Throwable {
    // GIVEN
    StageControl control = mock(StageControl.class);
    final AtomicLong nextId = new AtomicLong();
    RecordStore<RelationshipGroupRecord> store = mock(RecordStore.class);
    when(store.nextId()).thenAnswer(invocation -> nextId.incrementAndGet());
    doAnswer(invocation -> {
        invocation.getArgumentAt(0, RelationshipGroupRecord.class).setFirstOut(1);
        return null;
    }).when(store).prepareForCommit(any(RelationshipGroupRecord.class));
    Configuration config = Configuration.withBatchSize(DEFAULT, 10);
    EncodeGroupsStep encoder = new EncodeGroupsStep(control, config, store);
    // WHEN
    encoder.start(Step.ORDER_SEND_DOWNSTREAM);
    Catcher catcher = new Catcher();
    encoder.process(batch(new Group(1, 3), new Group(2, 3), new Group(3, 4)), catcher);
    encoder.process(batch(new Group(4, 2), new Group(5, 10)), catcher);
    encoder.process(batch(new Group(6, 35)), catcher);
    encoder.process(batch(new Group(7, 2)), catcher);
    encoder.endOfUpstream();
    while (!encoder.isCompleted()) {
        Thread.sleep(10);
    }
    encoder.close();
    // THEN
    assertEquals(4, catcher.batches.size());
    long lastOwningNodeLastBatch = -1;
    for (RelationshipGroupRecord[] batch : catcher.batches) {
        assertBatch(batch, lastOwningNodeLastBatch);
        lastOwningNodeLastBatch = batch[batch.length - 1].getOwningNode();
    }
}

Example 2 with StageControl

use of org.neo4j.unsafe.impl.batchimport.staging.StageControl in project neo4j by neo4j.

the class PropertyEncoderStepTest method shouldGrowPropertyBlocksArrayProperly.

@SuppressWarnings("unchecked")
@Test
public void shouldGrowPropertyBlocksArrayProperly() throws Exception {
    // GIVEN
    StageControl control = mock(StageControl.class);
    BatchingPropertyKeyTokenRepository tokens = new BatchingPropertyKeyTokenRepository(neoStores.getPropertyKeyTokenStore());
    Step<Batch<InputNode, NodeRecord>> step = new PropertyEncoderStep<>(control, DEFAULT, tokens, neoStores.getPropertyStore());
    @SuppressWarnings("rawtypes") Step downstream = mock(Step.class);
    step.setDownstream(downstream);
    // WHEN
    step.start(0);
    step.receive(0, smallbatch());
    step.endOfUpstream();
    awaitCompleted(step, control);
    // THEN
    verify(downstream).receive(anyLong(), any());
    verifyNoMoreInteractions(control);
    step.close();
}

Example 3 with StageControl

use of org.neo4j.unsafe.impl.batchimport.staging.StageControl in project neo4j by neo4j.

the class UpdateRecordsStepTest method ioThroughputStatDoesNotOverflow.

@Test
public void ioThroughputStatDoesNotOverflow() throws Throwable {
    // store with huge record size to force overflow and not create huge batch of records
    RecordStore<NodeRecord> store = mock(RecordStore.class);
    when(store.getRecordSize()).thenReturn(Integer.MAX_VALUE / 2);
    Configuration configuration = mock(Configuration.class);
    StageControl stageControl = mock(StageControl.class);
    UpdateRecordsStep<NodeRecord> step = new UpdateRecordsStep<>(stageControl, configuration, store);
    NodeRecord record = new NodeRecord(1);
    record.setInUse(true);
    NodeRecord[] batch = new NodeRecord[11];
    Arrays.fill(batch, record);
    step.process(batch, mock(BatchSender.class));
    Stat stat = step.stat(Keys.io_throughput);
    assertThat(stat.asLong(), greaterThan(0L));
}

Example 4 with StageControl

use of org.neo4j.unsafe.impl.batchimport.staging.StageControl in project neo4j by neo4j.

the class UpdateRecordsStepTest method recordWithReservedIdIsSkipped.

@Test
public void recordWithReservedIdIsSkipped() throws Throwable {
    RecordStore<NodeRecord> store = mock(NodeStore.class);
    StageControl stageControl = mock(StageControl.class);
    UpdateRecordsStep<NodeRecord> step = new UpdateRecordsStep<>(stageControl, Configuration.DEFAULT, store);
    NodeRecord node1 = new NodeRecord(1);
    node1.setInUse(true);
    NodeRecord node2 = new NodeRecord(2);
    node2.setInUse(true);
    NodeRecord nodeWithReservedId = new NodeRecord(IdGeneratorImpl.INTEGER_MINUS_ONE);
    NodeRecord[] batch = { node1, node2, nodeWithReservedId };
    step.process(batch, mock(BatchSender.class));
    verify(store).prepareForCommit(node1);
    verify(store).updateRecord(node1);
    verify(store).prepareForCommit(node2);
    verify(store).updateRecord(node2);
    verify(store, never()).prepareForCommit(nodeWithReservedId);
    verify(store, never()).updateRecord(nodeWithReservedId);
}