Skip to content

fix(pyamber): repair LinkedBlockingMultiQueue removal paths; add unit tests#6903

Merged
mengw15 merged 1 commit into
apache:mainfrom
mengw15:fix/6899-lbmq-remove-and-owner-binding
Jul 26, 2026
Merged

fix(pyamber): repair LinkedBlockingMultiQueue removal paths; add unit tests#6903
mengw15 merged 1 commit into
apache:mainfrom
mengw15:fix/6899-lbmq-remove-and-owner-binding

Conversation

@mengw15

@mengw15 mengw15 commented Jul 26, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this PR?

LinkedBlockingMultiQueue is the message queue behind every Python UDF worker: InternalQueue wraps it to give control messages their own priority lane ahead of data batches, and to implement pause and backpressure by disabling the data sub-queues.

It sat at ~70% line coverage. Writing tests for the uncovered parts surfaced two real defects, both on removal paths — which is exactly why those lines had never been exercised. This PR fixes both and adds the tests that pin the corrected behavior.

Neither defect is reachable from production code today. internal_queue.py is the only caller of this class and it uses put / get / size / enable / disable / is_empty / is_enabled / in_mem_size / add_sub_queue, so no existing caller changes behavior. They matter for the next caller: InternalQueue._queue_ids only ever grows, and whoever adds per-channel teardown will reach for remove_sub_queue first.

Defect 1 — SubQueue.remove corrupts the chain and desynchronises the count

head is a sentinel, so the candidate node is always trail.next, with trail as its predecessor. remove matched trail.item — the predecessor's own item — and then unlinked trail.next. unlink in turn cleared trail.item rather than the removed node's, and never decremented self.count.

On ["a", "b", "c"], remove("a") returns True and leaves:

before:  head(None) -> "a" -> "b" -> "c"     count=3  total_count=3
after:   head(None) -> None -> "c"           count=3  total_count=2
  • "a" is blanked in place instead of unlinked, and "b" is dropped
  • SubQueue.size() stays at 3 while the queue total goes to 2
  • the next get() returns None, a value put explicitly rejects
  • the last element can never match, because the loop stops before checking it — remove on a single-element sub-queue returns False

Fixed by matching trail.next.item, clearing next_.item, and decrementing self.count in unlink.

Defect 2 — remove_sub_queue raises AttributeError for any enabled sub-queue

The inner classes here carry an owner back-pointer to the outer queue, and @inner decides what it points at from how the inner class was reached:

  • self.PriorityGroup(...)owner is the queue object
  • LinkedBlockingMultiQueue.PriorityGroup(...)owner is the class

add_sub_queue used the second form. total_count is assigned in __init__, so it lives on the object, not on the class — and self.owner.total_count.get_and_dec(...) in remove_queue failed with:

AttributeError: type object 'LinkedBlockingMultiQueue' has no attribute 'total_count'

The throw lands mid-mutation, leaving the queue neither removed nor not-removed:

before:  sub_queues=['control','data']  groups=[(0,['control']),(1,['data'])]  total_count=3
after:   sub_queues=['control']         groups=[(0,['control']),(1,[])]        total_count=3

total_count still counts the removed sub-queue's items, and the now-empty PriorityGroup is never dropped — so the next get() raises IndexError on queues[self.next_idx] in get_next_sub_queue. Disabled sub-queues skip the failing line, which is the only reason this path ever appeared to work.

Fixed by constructing through self.PriorityGroup(...), matching how SubQueue is already built. Node and DefaultSubQueueSelection also use class access but never read self.owner, so they are left unchanged.

Tests

30 new tests cover peek at all three levels, SubQueue.clear, the corrected removal paths, PriorityGroup.remove_queue, remove_sub_queue, and the small guards. The queue fixture moved to module scope so the new classes can share it. Two current behaviors are pinned rather than changed: clear() does not reset in_mem_size, and SubQueue.__str__ only supports str items.

Any related issues, documentation, discussions?

Closes #6899.

How was this PR tested?

The target file passes with 42 tests (12 existing, 30 new), and the full pytest -m "not integration" suite passes with no regressions. ruff check and ruff format --check are clean.

Both defects were reproduced before the fix and confirmed gone after. To check the failure path, each fix was reverted in turn: reverting the remove fix reddens 4 removal tests, and reverting the owner-binding fix brings back the original AttributeError. The new tests genuinely guard both fixes.

Was this PR authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Fable 5)

… tests

SubQueue.remove matched the predecessor's item but unlinked the following
node, and unlink cleared the wrong item and never decremented SubQueue.count,
so removing an item dropped its successor and made the next get() return None.
add_sub_queue built PriorityGroup through class access, leaving owner bound to
the class so remove_sub_queue raised AttributeError for any enabled sub-queue.
@github-actions

Copy link
Copy Markdown
Contributor

Automated Reviewer Suggestions

Based on the git blame history of the changed files, we recommend the following reviewers:

  • Contributors with relevant context: @Yicong-Huang, @aglinxinyuan
    You can notify them by mentioning @Yicong-Huang, @aglinxinyuan in a comment.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR improves the correctness and test coverage of LinkedBlockingMultiQueue (PyAmber’s internal keyed multi-queue) by fixing two removal-path defects and adding a comprehensive unit-test suite that exercises previously uncovered behavior.

Changes:

  • Fix SubQueue.remove() / unlink() to correctly identify and unlink the target node, clear the removed node’s item, and keep SubQueue.count consistent with owner.total_count.
  • Fix add_sub_queue() to construct PriorityGroup via instance access (self.PriorityGroup(...)) so the @inner descriptor binds owner correctly, enabling remove_sub_queue() for enabled sub-queues.
  • Add ~30 new unit tests covering peek() at all levels, removal/clear paths, remove_sub_queue(), and various guard behaviors; refactor the shared queue fixture to module scope.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py Fix removal logic/count bookkeeping and correct PriorityGroup construction for proper owner binding.
amber/src/test/python/core/util/customized_queue/test_linked_blocking_multi_queue.py Expand unit test coverage to exercise peek/removal/clear and guard behaviors; share a module-level queue fixture.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.57%. Comparing base (4af3a85) to head (51751fe).

Additional details and impacted files
@@             Coverage Diff              @@
##               main    #6903      +/-   ##
============================================
+ Coverage     78.39%   78.57%   +0.18%     
  Complexity     3723     3723              
============================================
  Files          1161     1161              
  Lines         46066    46067       +1     
  Branches       5107     5107              
============================================
+ Hits          36112    36199      +87     
+ Misses         8343     8257      -86     
  Partials       1611     1611              
Flag Coverage Δ *Carryforward flag
access-control-service 70.00% <ø> (ø) Carriedforward from 4af3a85
agent-service 76.76% <ø> (ø) Carriedforward from 4af3a85
amber 71.51% <ø> (ø) Carriedforward from 4af3a85
computing-unit-managing-service 20.49% <ø> (ø) Carriedforward from 4af3a85
config-service 66.66% <ø> (ø) Carriedforward from 4af3a85
file-service 67.21% <ø> (ø) Carriedforward from 4af3a85
frontend 82.64% <ø> (ø) Carriedforward from 4af3a85
notebook-migration-service 78.94% <ø> (ø) Carriedforward from 4af3a85
pyamber 95.40% <100.00%> (+1.91%) ⬆️
workflow-compiling-service 55.14% <ø> (ø) Carriedforward from 4af3a85

*This pull request uses carry forward flags. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@github-actions

Copy link
Copy Markdown
Contributor

⚠️ Benchmark changes need a look

🟢 0 better · 🔴 10 worse · ⚪ 5 noise (<±5%) · 0 without baseline

Compared against main 4af3a85 benchmarked on this same runner, so the delta is largely free of cross-runner hardware noise. The "7d avg" column still reflects the gh-pages dashboard. Treat <±5% as noise unless repeated.

Dashboard · Run

config throughput MB/s latency max Δ latest / 7d
🔴 bs=10 sw=10 sl=64 363 0.221 26,792/38,040/38,040 us 🔴 -15.3% / 🔴 +129.2%
🔴 bs=100 sw=10 sl=64 776 0.473 128,596/161,670/161,670 us 🔴 +21.6% / 🔴 +48.2%
bs=1000 sw=10 sl=64 904 0.552 1,105,721/1,196,679/1,196,679 us ⚪ within ±5% / 🔴 +14.4%
Baseline details

Latest main 4af3a85 from same runner

config metric PR latest main 7d avg Δ latest Δ 7d
bs=10 sw=10 sl=64 throughput 363 tuples/sec 427 tuples/sec 754.55 tuples/sec -15.0% -51.9%
bs=10 sw=10 sl=64 MB/s 0.221 MB/s 0.261 MB/s 0.461 MB/s -15.3% -52.0%
bs=10 sw=10 sl=64 p50 26,792 us 23,372 us 12,816 us +14.6% +109.1%
bs=10 sw=10 sl=64 p95 38,040 us 33,420 us 16,594 us +13.8% +129.2%
bs=10 sw=10 sl=64 p99 38,040 us 33,420 us 19,806 us +13.8% +92.1%
bs=100 sw=10 sl=64 throughput 776 tuples/sec 822 tuples/sec 969.38 tuples/sec -5.6% -19.9%
bs=100 sw=10 sl=64 MB/s 0.473 MB/s 0.502 MB/s 0.592 MB/s -5.8% -20.1%
bs=100 sw=10 sl=64 p50 128,596 us 121,866 us 103,584 us +5.5% +24.1%
bs=100 sw=10 sl=64 p95 161,670 us 133,000 us 109,097 us +21.6% +48.2%
bs=100 sw=10 sl=64 p99 161,670 us 133,000 us 117,304 us +21.6% +37.8%
bs=1000 sw=10 sl=64 throughput 904 tuples/sec 915 tuples/sec 1,004 tuples/sec -1.2% -9.9%
bs=1000 sw=10 sl=64 MB/s 0.552 MB/s 0.559 MB/s 0.613 MB/s -1.3% -9.9%
bs=1000 sw=10 sl=64 p50 1,105,721 us 1,092,994 us 1,002,357 us +1.2% +10.3%
bs=1000 sw=10 sl=64 p95 1,196,679 us 1,169,952 us 1,046,463 us +2.3% +14.4%
bs=1000 sw=10 sl=64 p99 1,196,679 us 1,169,952 us 1,073,661 us +2.3% +11.5%
Raw CSV
config_idx,batch_size,schema_width,string_len,num_batches,total_ms,total_tuples,total_bytes,tuples_per_sec,mb_per_sec,lat_p50_us,lat_p95_us,lat_p99_us
0,10,10,64,20,551.27,200,128000,363,0.221,26792.21,38039.97,38039.97
1,100,10,64,20,2578.96,2000,1280000,776,0.473,128595.50,161670.25,161670.25
2,1000,10,64,20,22112.32,20000,12800000,904,0.552,1105720.77,1196678.75,1196678.75

@aglinxinyuan aglinxinyuan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@mengw15
mengw15 added this pull request to the merge queue Jul 26, 2026
Merged via the queue into apache:main with commit 38d94ea Jul 26, 2026
24 checks passed
@mengw15
mengw15 deleted the fix/6899-lbmq-remove-and-owner-binding branch July 26, 2026 06:14
@chenlica

Copy link
Copy Markdown
Contributor

@mengw15 I found the description very hard to read. Can you improve it?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Extend pyamber LinkedBlockingMultiQueue unit test coverage

5 participants