Fix task bootstrapping & simplify segment load/drop flows by abhishekrb19 · Pull Request #16475 · apache/druid

abhishekrb19 · 2024-05-20T18:33:58Z

Problems

The segment bootstrapping logic relies on the presence of configured storage locations. That is, only servers that have storage location configured will act like a "segment data server". Tasks on the other hand don't have this static configuration because tasks can dynamically run on any slot. The CliPeon injects a list of dynamic StorageLocations. This List<StorageLocation> is not currently present in SegmentLoadDropHandler, which does the bootstrapping of segments as part of its start() lifecycle. Instead, this class directly uses config.getLocations(), which isn't set by the TaskToolboxFactory when it instantiates SegmentCacheManager.
One of the overloaded constructors marked @VisibleForTesting in SegmentLocalCacheManager is also being used by tasks via SegmentCacheManagerFactory. This "test constructor" also does not initialize the thread pool executors, so segments won't be loaded into the page cache even if loader config has the appropriate settings.

Core changes

The core part of the fix in SegmentLoadDropHandler is to rely on the source of truth for the caching layer: SegmentCacheManager. Currently, the various interactions of the classes are complex and have abstraction leaks in its implementation:

SegmentManager -> SegmentCacheLoader -> SegmentCacheManager
SegmentLoadDropHandler -> SegmentManager -> SegmentCacheLoader -> SegmentCacheManager
SegmentLoadDropHandler -> SegmentCacheManager

The SegmentLoadDropHandler directly handles the segment info file creation and deletion. These low-level file operations should really be owned and handled by the SegmentCacheManager and the classes that interact with should just operate at the granularity of Segment objects.

With this patch, the updated interactions are as follows:

SegmentManager -> SegmentCacheManager
SegmentLoadDropHandler -> SegmentManager -> SegmentCacheManager

The addition and deletion of segment info files are handled by the SegmentLocalCacheManager class itself. The SegmentLoadDropHandler gets the set of cached segments from the SegmentManager. To that effect, the dependencies in SegmentLoadDropHandler have been cleaned up, and this class doesn't fiddle with the segment files directly.

In order to simplify this, we did the following:

Removed the SegmentLoader interface and its concrete implementation SegmentLocalCacheLoader. The implementation was acting as a simple factory of Segment objects from the underlying segment files, which has now been moved to SegmentLocalCacheManager, still providing convenience for any test overrides for the factory.
Added the following methods to SegmentCacheManager for callers to interact with:
- boolean canHandleSegments(): A high-level method that informs callers whether the caching layer can handle segments or not.
- List getCachedSegments()
- void storeInfoFile(DataSegment segment)
- void removeInfoFile(DataSegment segment)
- ReferenceCountingSegment getSegment(DataSegment segment)
- ReferenceCountingSegment getBootstrapSegment(DataSegment segment, SegmentLazyLoadFailCallback loadFailed)
- void loadSegmentIntoPageCache(DataSegment segment)
- void loadSegmentIntoPageCacheOnBootstrap(DataSegment segment)

SegmentManager is now responsible for calling storeInfoFile(segment) once the segment bookkeeping stuff is successful as part of the load segment flow.
The bootstrap executor was initialized by SegmentLoadDropHandler, which is now moved to the SegmentLocalCacheManager., along with the initialization of the segment download executor.
Methods in SegmentManager and SegmentLocalCacheManager have bene broken down into loadSegment() and loadBootstrapSegment() variants, because segment lazy loading and the fail callback are only applicable for the bootstrap segment flow. This cleanup ensures that callers don't pass unnecessary null arguments anymore and makes the intent clearer.
Remove some VisibleForTesting hooks, including the problematic "test constructor" in SegmentLocalCacheManager noted in the problems section. Where possible, the test code uses the same initialization logic as the core code.

Test changes

Added new tests in SegmentLocalCacheManagerTest and SegmentManagerTest.
Moved the existing tests from SegmentLoaderTest to SegmentLocalCacheManagerTest.
Added TestSegmentUtils for commonly used test utilities that has test implementation support for loading DataSegment, Segment factory, etc.

Future work

We can further simplify and cleanup these interactions by moving the functionality of SegmentLoadDropHandler to SegmentManager such that the announcements moves to the manager layer itself. That'll entail a slightly larger refactor.
Cleanup the test code in SegmentLocalCacheManagerTest

Both of these require moderate to large refactoring changes, which can be tackled separately.

Release note

Fixed an issue in task bootstrapping that prevented tasks from accepting any segment assignments, including broadcast segments.

This PR has:

been self-reviewed.
a release note entry in the PR description.
added Javadocs for most classes and all non-trivial methods. Linked related entities via Javadoc links.
added unit tests or modified existing tests to cover new code paths, ensuring the threshold for code coverage is met.
been tested in a test Druid cluster.

- The load drop handler code talks to the local cache manager via SegmentManager.

+  public void storeInfoFile(DataSegment segment) throws IOException
+  {
+    final File segmentInfoCacheFile = new File(getInfoDir(), segment.getId().toString());
+    if (!segmentInfoCacheFile.exists()) {


+    final File factoryJson = new File(segmentFiles, "factory.json");
+    final SegmentizerFactory factory;
+
+    if (factoryJson.exists()) {


+
+    if (factoryJson.exists()) {
+      try {
+        factory = jsonMapper.readValue(factoryJson, SegmentizerFactory.class);


+      try {
+        byte[] bytes = new byte[size];
+        ThreadLocalRandom.current().nextBytes(bytes);
+        Files.write(bytes, segmentFile);


+        byte[] bytes = new byte[size];
+        ThreadLocalRandom.current().nextBytes(bytes);
+        Files.write(bytes, segmentFile);
+        Files.write("{\"type\":\"testSegmentFactory\"}".getBytes(StandardCharsets.UTF_8), factoryJson);


kfaraz · 2024-05-21T06:13:05Z

@abhishekrb19 , thanks for the detailed description.

IIUC, this PR is basically doing the following:

Inject the correct value of storage locations into SegmentLoadDropHandler
Refactor and cleanup classes

Is there any other bug fix or feature included in this PR too?
Just wanted to confirm before I start the actual review.

abhishekrb19 · 2024-05-21T07:09:29Z

IIUC, this PR is basically doing the following:

Inject the correct value of storage locations into SegmentLoadDropHandler

Refactor and cleanup classes

That's about right, except that we're not injecting storage locations into SegmentLoadDropHandler to fix the bug. Instead, SegmentLoadDropHandler relies on SegmentManager and SegmentCacheManager to provide the needed information. With this patch, SegmentCacheManager essentially serves as the source of truth for low-level segment info file operations, and this class already has the list of storage locations injected correctly.

Is there any other bug fix or feature included in this PR too? Just wanted to confirm before I start the actual review.

There's no other bug fix or feature in this PR. I do have another related change that touches some of this code, but given the size of this PR, I will add it separately to make reviewing easier. 🙂

kfaraz · 2024-05-21T08:21:09Z

Makes sense. Thanks for the clarification, @abhishekrb19 !

…e usages.

The intellij-inspect tool doesn't seem to correctly inspect lambda usages. See ScheduledExecutors.

…/incubator-druid into fixup_task_bootstrap_inject

Fix task bootstrap locations.

388e100

github-actions Bot added Area - Batch Ingestion Area - Ingestion Area - MSQ For multi stage queries - https://github.com/apache/druid/issues/12262 labels May 20, 2024

abhishekrb19 added 2 commits May 20, 2024 11:40

Remove dependency of SegmentCacheManager from SegmentLoadDropHandler.

0cd0f5a

- The load drop handler code talks to the local cache manager via SegmentManager.

Clean up unused imports and stuff.

6ddcdd5

github-advanced-security AI found potential problems May 20, 2024

View reviewed changes

Test fixes.

e91b82c

github-actions Bot added the Area - Querying label May 20, 2024

abhishekrb19 force-pushed the fixup_task_bootstrap_inject branch from be28082 to aece4e6 Compare May 20, 2024 21:51

Intellij inspections and test bind.

6e563aa

abhishekrb19 force-pushed the fixup_task_bootstrap_inject branch from aece4e6 to 6e563aa Compare May 21, 2024 00:26

Clean up dependencies some more

b242021

abhishekrb19 requested review from cheddar and kfaraz May 21, 2024 01:31

github-advanced-security AI found potential problems May 21, 2024

View reviewed changes

Comment thread server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java Fixed

abhishekrb19 added 2 commits May 20, 2024 20:12

Extract test load spec and factory to its own class.

9666512

Cleanup test util

e2feca5

github-advanced-security AI found potential problems May 21, 2024

View reviewed changes

abhishekrb19 added 2 commits May 20, 2024 22:56

Pull SegmentForTesting out to TestSegmentUtils.

fb018bb

Fix up.

a51d92a

Minor changes to infoDir

e2b6f18

abhishekrb19 added 5 commits May 21, 2024 10:02

Merge branch 'master' into fixup_task_bootstrap_inject

633cc36

Replace server announcer mock and verify that.

92a7bac

Add tests.

ee74c67

Merge branch 'master' into fixup_task_bootstrap_inject

302834d

Update javadocs.

8ad46b9

abhishekrb19 added 3 commits May 28, 2024 08:57

Remove the other constructor that wasn't initializing fully and updat…

9bf4ace

…e usages.

Cleanup withInfoDir() builder and unnecessary test hooks.

c6b8cc8

Remove mocks and elaborate on comments.

16658e1

abhishekrb19 force-pushed the fixup_task_bootstrap_inject branch from 19c00e4 to 16658e1 Compare May 29, 2024 06:08

abhishekrb19 added 4 commits May 28, 2024 23:18

Merge branch 'master' into fixup_task_bootstrap_inject

ae49b07

Commentary

e162311

Merge branch 'master' into fixup_task_bootstrap_inject

b6f8f20

Fix a few Intellij inspection warnings.

e02f5cb

abhishekrb19 force-pushed the fixup_task_bootstrap_inject branch from 1e028cb to e02f5cb Compare May 29, 2024 19:49

Suppress corePoolSize intellij-inspect warning.

98db446

The intellij-inspect tool doesn't seem to correctly inspect lambda usages. See ScheduledExecutors.

abhishekrb19 requested a review from jon-wei May 30, 2024 21:37

imply-cheddar approved these changes May 31, 2024

View reviewed changes

jon-wei approved these changes Jun 3, 2024

View reviewed changes

abhishekrb19 added 3 commits June 3, 2024 08:38

Update docs and add more tests.

8024d51

Merge branch 'master' into fixup_task_bootstrap_inject

da73aa6

Use hamcrest for asserting order on expectation.

f11da12

abhishekrb19 force-pushed the fixup_task_bootstrap_inject branch from d2d1c2d to f11da12 Compare June 3, 2024 16:37

abhishekrb19 added the Release Notes label Jun 3, 2024

abhishekrb19 changed the title ~~Fix task bootstrapping.~~ Fix task bootstrapping & simplify segment load/drop flows Jun 3, 2024

Merge branch 'master' into fixup_task_bootstrap_inject

b036808

asdf2014 approved these changes Jun 4, 2024

View reviewed changes

abhishekrb19 added 3 commits June 4, 2024 07:09

Shutdown bootstrap exec.

c2c9fd1

Merge branch 'fixup_task_bootstrap_inject' of github.com:abhishekrb19…

a16d69e

…/incubator-druid into fixup_task_bootstrap_inject

Fix checkstyle

f97d10f

abhishekrb19 force-pushed the fixup_task_bootstrap_inject branch from c262ed4 to f97d10f Compare June 4, 2024 14:38

abhishekrb19 merged commit b9ba286 into apache:master Jun 4, 2024

abhishekrb19 mentioned this pull request Jun 21, 2024

Support for bootstrap segments #16609

Merged

10 tasks

abhishekrb19 deleted the fixup_task_bootstrap_inject branch June 24, 2024 15:39

kfaraz added this to the 31.0.0 milestone Oct 4, 2024

kfaraz mentioned this pull request Oct 11, 2024

[DRAFT] 31.0.0 Release Notes #17332

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix task bootstrapping & simplify segment load/drop flows#16475

Fix task bootstrapping & simplify segment load/drop flows#16475
abhishekrb19 merged 48 commits intoapache:masterfrom
abhishekrb19:fixup_task_bootstrap_inject

abhishekrb19 commented May 20, 2024 •

edited

Loading

Uh oh!

Check failure

Uh oh!

Check failure

Check failure

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Check failure

Check failure

kfaraz commented May 21, 2024 •

edited

Loading

Uh oh!

abhishekrb19 commented May 21, 2024 •

edited

Loading

Uh oh!

kfaraz commented May 21, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

abhishekrb19 commented May 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problems

Core changes

With this patch, the updated interactions are as follows:

Test changes

Future work

Release note

Uh oh!

Check failure

Uh oh!

Check failure

Check failure

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Check failure

Check failure

kfaraz commented May 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

abhishekrb19 commented May 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kfaraz commented May 21, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

abhishekrb19 commented May 20, 2024 •

edited

Loading

kfaraz commented May 21, 2024 •

edited

Loading

abhishekrb19 commented May 21, 2024 •

edited

Loading