github-advanced-security[bot] commented on code in PR #17653: URL: https://github.com/apache/druid/pull/17653#discussion_r1931567954
########## server/src/main/java/org/apache/druid/metadata/segment/cache/DatasourceSegmentCache.java: ########## @@ -0,0 +1,596 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.metadata.segment.cache; + +import org.apache.druid.error.DruidException; +import org.apache.druid.java.util.common.CloseableIterators; +import org.apache.druid.java.util.common.parsers.CloseableIterator; +import org.apache.druid.metadata.PendingSegmentRecord; +import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec; +import org.apache.druid.server.http.DataSegmentPlus; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.SegmentTimeline; +import org.joda.time.DateTime; +import org.joda.time.Interval; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.function.Predicate; +import java.util.stream.Collectors; + +/** + * Datasource-level cache for segments and pending segments. + */ +class DatasourceSegmentCache extends BaseCache +{ + private final String dataSource; + + /** + * Used to obtain the segment for a given ID so that it can be updated in the + * timeline. + */ + private final Map<String, DataSegmentPlus> idToUsedSegment = new HashMap<>(); + + /** + * Current state of segments as seen by the cache. + */ + private final Map<String, SegmentState> idToSegmentState = new HashMap<>(); + + /** + * Allows lookup of visible segments for a given interval. + */ + private final SegmentTimeline usedSegmentTimeline = SegmentTimeline.forSegments(Set.of()); + + private final Map<Interval, Map<String, PendingSegmentRecord>> + intervalToPendingSegments = new HashMap<>(); + + private final Map<Interval, Map<String, Integer>> + intervalVersionToHighestUnusedPartitionNumber = new HashMap<>(); + + DatasourceSegmentCache(String dataSource) + { + super(true); + this.dataSource = dataSource; + } + + void clear() + { + withWriteLock(() -> { + idToSegmentState.clear(); + idToUsedSegment.clear(); + intervalVersionToHighestUnusedPartitionNumber.clear(); + idToUsedSegment.values().forEach(s -> usedSegmentTimeline.remove(s.getDataSegment())); + }); + } + + boolean isEmpty() + { + return withReadLock(() -> idToSegmentState.isEmpty() && intervalToPendingSegments.isEmpty()); + } + + /** + * Checks if a segment needs to be refreshed. A refresh is required if the + * cache has no known state for the given segment or if the metadata store + * has a more recent last_updated_time than the cache. + */ + boolean shouldRefreshSegment(String segmentId, SegmentState metadataState) + { + return withReadLock(() -> { + final SegmentState cachedState = idToSegmentState.get(segmentId); + return cachedState == null + || cachedState.getLastUpdatedTime().isBefore(metadataState.getLastUpdatedTime()); + }); + } + + /** + * Checks if a pending segment needs to be refreshed in the cache. + */ + boolean shouldRefreshPendingSegment(PendingSegmentRecord record) + { + final SegmentIdWithShardSpec segmentId = record.getId(); + return withReadLock( + () -> intervalToPendingSegments.getOrDefault(segmentId.getInterval(), Map.of()) + .containsKey(segmentId.toString()) + ); + } + + boolean refreshUnusedSegment(String segmentId, SegmentState newState) + { + if (newState.isUsed()) { + return false; + } + + return withWriteLock(() -> { + if (!shouldRefreshSegment(segmentId, newState)) { + return false; + } + + final SegmentState oldState = idToSegmentState.put(segmentId, newState); + + if (oldState != null && oldState.isUsed()) { + // Segment has transitioned from used to unused + DataSegmentPlus segment = idToUsedSegment.remove(segmentId); + if (segment != null) { + usedSegmentTimeline.remove(segment.getDataSegment()); + } + } + + addUnusedSegmentId(segmentId); + return true; + }); + } + + boolean refreshUsedSegment(DataSegmentPlus segmentPlus) + { + final DataSegment segment = segmentPlus.getDataSegment(); + final String segmentId = getId(segment); + + final SegmentState newState = new SegmentState( + Boolean.TRUE.equals(segmentPlus.getUsed()), + segmentPlus.getUsedStatusLastUpdatedDate() + ); + if (!newState.isUsed()) { + return refreshUnusedSegment(segmentId, newState); + } + + return withWriteLock(() -> { + if (!shouldRefreshSegment(segmentId, newState)) { + return false; + } + + final SegmentState oldState = idToSegmentState.put(segmentId, newState); + final DataSegmentPlus oldSegmentPlus = idToUsedSegment.put(segmentId, segmentPlus); + + if (oldState == null) { + // This is a new segment + } else if (oldState.isUsed()) { + // Segment payload may have changed + if (oldSegmentPlus != null) { + usedSegmentTimeline.remove(oldSegmentPlus.getDataSegment()); + } + } else { + // Segment has transitioned from unused to used + removeUnusedSegmentId(segmentId); + } + + usedSegmentTimeline.add(segment); + return true; + }); + } + + int removeSegmentIds(Set<String> segmentIds) + { + return withWriteLock(() -> { + int removedCount = 0; + for (String segmentId : segmentIds) { + SegmentState state = idToSegmentState.remove(segmentId); + if (state != null) { + ++removedCount; + } + + removeUnusedSegmentId(segmentId); + + final DataSegmentPlus segment = idToUsedSegment.remove(segmentId); + if (segment != null) { + usedSegmentTimeline.remove(segment.getDataSegment()); + } + } + + return removedCount; + }); + } + + void resetMaxUnusedIds(Map<Interval, Map<String, Integer>> intervalVersionToHighestPartitionNumber) + { + withWriteLock(() -> { + this.intervalVersionToHighestUnusedPartitionNumber.clear(); + this.intervalVersionToHighestUnusedPartitionNumber.putAll(intervalVersionToHighestPartitionNumber); + }); + } + + private void addUnusedSegmentId(String id) + { + final SegmentId segmentId = SegmentId.tryParse(dataSource, id); + if (segmentId == null) { + return; + } + + final int partitionNum = segmentId.getPartitionNum(); + intervalVersionToHighestUnusedPartitionNumber + .computeIfAbsent(segmentId.getInterval(), i -> new HashMap<>()) + .merge(segmentId.getVersion(), partitionNum, Math::max); + } + + private void removeUnusedSegmentId(String segmentId) Review Comment: ## Useless parameter The parameter 'segmentId' is never used. [Show more details](https://github.com/apache/druid/security/code-scanning/8658) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
