Examples with SegmentLock org.apache.druid.indexing.common.SegmentLock used on opensource projects

Example 1 with SegmentLock

use of org.apache.druid.indexing.common.SegmentLock in project druid by druid-io.

the class TaskLocks method findLocksForSegments.

public static List<TaskLock> findLocksForSegments(final Task task, final TaskLockbox taskLockbox, final Collection<DataSegment> segments) {
    final NavigableMap<DateTime, List<TaskLock>> taskLockMap = getTaskLockMap(taskLockbox, task);
    if (taskLockMap.isEmpty()) {
        return Collections.emptyList();
    }
    final List<TaskLock> found = new ArrayList<>();
    segments.forEach(segment -> {
        final Entry<DateTime, List<TaskLock>> entry = taskLockMap.floorEntry(segment.getInterval().getStart());
        if (entry == null) {
            throw new ISE("Can't find lock for the interval of segment[%s]", segment.getId());
        }
        final List<TaskLock> locks = entry.getValue();
        locks.forEach(lock -> {
            if (lock.getGranularity() == LockGranularity.TIME_CHUNK) {
                final TimeChunkLock timeChunkLock = (TimeChunkLock) lock;
                if (timeChunkLock.getInterval().contains(segment.getInterval()) && timeChunkLock.getDataSource().equals(segment.getDataSource()) && timeChunkLock.getVersion().compareTo(segment.getVersion()) >= 0) {
                    found.add(lock);
                }
            } else {
                final SegmentLock segmentLock = (SegmentLock) lock;
                if (segmentLock.getInterval().contains(segment.getInterval()) && segmentLock.getDataSource().equals(segment.getDataSource()) && segmentLock.getVersion().compareTo(segment.getVersion()) >= 0 && segmentLock.getPartitionId() == segment.getShardSpec().getPartitionNum()) {
                    found.add(lock);
                }
            }
        });
    });
    return found;
}

Example 2 with SegmentLock

use of org.apache.druid.indexing.common.SegmentLock in project druid by druid-io.

the class SegmentAllocateActionTest method testResumeSequence.

@Test
public void testResumeSequence() {
    final Task task = NoopTask.create();
    taskActionTestKit.getTaskLockbox().add(task);
    final Map<Integer, SegmentIdWithShardSpec> allocatedPartyTimeIds = new HashMap<>();
    final Map<Integer, SegmentIdWithShardSpec> allocatedFutureIds = new HashMap<>();
    final SegmentIdWithShardSpec id1 = allocate(task, PARTY_TIME, Granularities.NONE, Granularities.HOUR, "s1", null);
    Assert.assertNotNull(id1);
    allocatedPartyTimeIds.put(id1.getShardSpec().getPartitionNum(), id1);
    final SegmentIdWithShardSpec id2 = allocate(task, THE_DISTANT_FUTURE, Granularities.NONE, Granularities.HOUR, "s1", id1.toString());
    Assert.assertNotNull(id2);
    allocatedFutureIds.put(id2.getShardSpec().getPartitionNum(), id2);
    final SegmentIdWithShardSpec id3 = allocate(task, PARTY_TIME, Granularities.NONE, Granularities.HOUR, "s1", id2.toString());
    Assert.assertNotNull(id3);
    allocatedPartyTimeIds.put(id3.getShardSpec().getPartitionNum(), id3);
    final SegmentIdWithShardSpec id4 = allocate(task, PARTY_TIME, Granularities.NONE, Granularities.HOUR, "s1", id1.toString());
    Assert.assertNull(id4);
    final SegmentIdWithShardSpec id5 = allocate(task, THE_DISTANT_FUTURE, Granularities.NONE, Granularities.HOUR, "s1", id1.toString());
    Assert.assertNotNull(id5);
    allocatedFutureIds.put(id5.getShardSpec().getPartitionNum(), id5);
    final SegmentIdWithShardSpec id6 = allocate(task, THE_DISTANT_FUTURE, Granularities.NONE, Granularities.MINUTE, "s1", id1.toString());
    Assert.assertNull(id6);
    final SegmentIdWithShardSpec id7 = allocate(task, THE_DISTANT_FUTURE, Granularities.NONE, Granularities.DAY, "s1", id1.toString());
    Assert.assertNotNull(id7);
    allocatedFutureIds.put(id7.getShardSpec().getPartitionNum(), id7);
    if (lockGranularity == LockGranularity.TIME_CHUNK) {
        final TaskLock partyLock = Iterables.getOnlyElement(FluentIterable.from(taskActionTestKit.getTaskLockbox().findLocksForTask(task)).filter(new Predicate<TaskLock>() {

            @Override
            public boolean apply(TaskLock input) {
                return input.getInterval().contains(PARTY_TIME);
            }
        }));
        final TaskLock futureLock = Iterables.getOnlyElement(FluentIterable.from(taskActionTestKit.getTaskLockbox().findLocksForTask(task)).filter(new Predicate<TaskLock>() {

            @Override
            public boolean apply(TaskLock input) {
                return input.getInterval().contains(THE_DISTANT_FUTURE);
            }
        }));
        assertSameIdentifier(id1, new SegmentIdWithShardSpec(DATA_SOURCE, Granularities.HOUR.bucket(PARTY_TIME), partyLock.getVersion(), new NumberedShardSpec(0, 0)));
        assertSameIdentifier(id2, new SegmentIdWithShardSpec(DATA_SOURCE, Granularities.HOUR.bucket(THE_DISTANT_FUTURE), futureLock.getVersion(), new NumberedShardSpec(0, 0)));
        assertSameIdentifier(id3, new SegmentIdWithShardSpec(DATA_SOURCE, Granularities.HOUR.bucket(PARTY_TIME), partyLock.getVersion(), new NumberedShardSpec(1, 0)));
    } else {
        final List<TaskLock> partyLocks = taskActionTestKit.getTaskLockbox().findLocksForTask(task).stream().filter(input -> input.getInterval().contains(PARTY_TIME)).collect(Collectors.toList());
        Assert.assertEquals(2, partyLocks.size());
        final Map<Integer, SegmentLock> partitionIdToLock = new HashMap<>();
        partyLocks.forEach(lock -> {
            Assert.assertEquals(LockGranularity.SEGMENT, lock.getGranularity());
            final SegmentLock segmentLock = (SegmentLock) lock;
            partitionIdToLock.put(segmentLock.getPartitionId(), segmentLock);
        });
        for (Entry<Integer, SegmentLock> entry : partitionIdToLock.entrySet()) {
            assertSameIdentifier(new SegmentIdWithShardSpec(DATA_SOURCE, Granularities.HOUR.bucket(PARTY_TIME), allocatedPartyTimeIds.get(entry.getKey()).getVersion(), new NumberedShardSpec(entry.getValue().getPartitionId(), 0)), allocatedPartyTimeIds.get(entry.getKey()));
        }
        final List<TaskLock> futureLocks = taskActionTestKit.getTaskLockbox().findLocksForTask(task).stream().filter(input -> input.getInterval().contains(THE_DISTANT_FUTURE)).collect(Collectors.toList());
        Assert.assertEquals(1, futureLocks.size());
        partitionIdToLock.clear();
        futureLocks.forEach(lock -> {
            Assert.assertEquals(LockGranularity.SEGMENT, lock.getGranularity());
            final SegmentLock segmentLock = (SegmentLock) lock;
            partitionIdToLock.put(segmentLock.getPartitionId(), segmentLock);
        });
        for (Entry<Integer, SegmentLock> entry : partitionIdToLock.entrySet()) {
            assertSameIdentifier(new SegmentIdWithShardSpec(DATA_SOURCE, Granularities.HOUR.bucket(THE_DISTANT_FUTURE), allocatedFutureIds.get(entry.getKey()).getVersion(), new NumberedShardSpec(entry.getValue().getPartitionId(), 0)), allocatedFutureIds.get(entry.getKey()));
        }
    }
    Assert.assertNull(id4);
    assertSameIdentifier(id2, id5);
    Assert.assertNull(id6);
    assertSameIdentifier(id2, id7);
}

Example 3 with SegmentLock

use of org.apache.druid.indexing.common.SegmentLock in project druid by druid-io.

the class TaskLockboxTest method testSegmentLock.

@Test
public void testSegmentLock() throws InterruptedException {
    final Task task = NoopTask.create();
    lockbox.add(task);
    final LockResult lockResult = lockbox.lock(task, new SpecificSegmentLockRequest(TaskLockType.EXCLUSIVE, task, Intervals.of("2015-01-01/2015-01-02"), "v1", 3));
    Assert.assertTrue(lockResult.isOk());
    Assert.assertNull(lockResult.getNewSegmentId());
    Assert.assertTrue(lockResult.getTaskLock() instanceof SegmentLock);
    final SegmentLock segmentLock = (SegmentLock) lockResult.getTaskLock();
    Assert.assertEquals(TaskLockType.EXCLUSIVE, segmentLock.getType());
    Assert.assertEquals(task.getGroupId(), segmentLock.getGroupId());
    Assert.assertEquals(task.getDataSource(), segmentLock.getDataSource());
    Assert.assertEquals(Intervals.of("2015-01-01/2015-01-02"), segmentLock.getInterval());
    Assert.assertEquals("v1", segmentLock.getVersion());
    Assert.assertEquals(3, segmentLock.getPartitionId());
    Assert.assertEquals(task.getPriority(), segmentLock.getPriority().intValue());
    Assert.assertFalse(segmentLock.isRevoked());
}

Example 4 with SegmentLock

use of org.apache.druid.indexing.common.SegmentLock in project druid by druid-io.

the class TaskLockboxTest method testRequestForNewSegmentWithSegmentLock.

@Test
public void testRequestForNewSegmentWithSegmentLock() {
    final Task task = NoopTask.create();
    lockbox.add(task);
    allocateSegmentsAndAssert(task, "seq", 3, NumberedPartialShardSpec.instance());
    allocateSegmentsAndAssert(task, "seq2", 2, new NumberedOverwritePartialShardSpec(0, 3, (short) 1));
    final List<TaskLock> locks = lockbox.findLocksForTask(task);
    Assert.assertEquals(5, locks.size());
    int expectedPartitionId = 0;
    for (TaskLock lock : locks) {
        Assert.assertTrue(lock instanceof SegmentLock);
        final SegmentLock segmentLock = (SegmentLock) lock;
        Assert.assertEquals(expectedPartitionId++, segmentLock.getPartitionId());
        if (expectedPartitionId == 3) {
            expectedPartitionId = PartitionIds.NON_ROOT_GEN_START_PARTITION_ID;
        }
    }
}

Example 5 with SegmentLock

use of org.apache.druid.indexing.common.SegmentLock in project druid by druid-io.

the class TaskLockboxTest method testGetTimeChunkAndSegmentLockForSameGroup.

@Test
public void testGetTimeChunkAndSegmentLockForSameGroup() {
    final Task task1 = NoopTask.withGroupId("groupId");
    final Task task2 = NoopTask.withGroupId("groupId");
    lockbox.add(task1);
    lockbox.add(task2);
    Assert.assertTrue(lockbox.tryLock(task1, new TimeChunkLockRequest(TaskLockType.EXCLUSIVE, task1, Intervals.of("2017/2018"), null)).isOk());
    Assert.assertTrue(lockbox.tryLock(task2, new SpecificSegmentLockRequest(TaskLockType.EXCLUSIVE, task2, Intervals.of("2017/2018"), "version", 0)).isOk());
    final List<TaskLockPosse> posses = lockbox.getAllLocks().get(task1.getDataSource()).get(DateTimes.of("2017")).get(Intervals.of("2017/2018"));
    Assert.assertEquals(2, posses.size());
    Assert.assertEquals(LockGranularity.TIME_CHUNK, posses.get(0).getTaskLock().getGranularity());
    final TimeChunkLock timeChunkLock = (TimeChunkLock) posses.get(0).getTaskLock();
    Assert.assertEquals("none", timeChunkLock.getDataSource());
    Assert.assertEquals("groupId", timeChunkLock.getGroupId());
    Assert.assertEquals(Intervals.of("2017/2018"), timeChunkLock.getInterval());
    Assert.assertEquals(LockGranularity.SEGMENT, posses.get(1).getTaskLock().getGranularity());
    final SegmentLock segmentLock = (SegmentLock) posses.get(1).getTaskLock();
    Assert.assertEquals("none", segmentLock.getDataSource());
    Assert.assertEquals("groupId", segmentLock.getGroupId());
    Assert.assertEquals(Intervals.of("2017/2018"), segmentLock.getInterval());
    Assert.assertEquals(0, segmentLock.getPartitionId());
}