Kalvin2077 commented on PR #3809:
URL: https://github.com/apache/celeborn/pull/3809#issuecomment-5600558119
### **Ask** ❓
test
### **Answer:**
The PR adds meaningful regression tests in `TierWriterSuite`, and they are
focused on the right problem: making sure `TierWriter` does not corrupt
memory/disk accounting or leak `ByteBuf` references when
`CompositeByteBuf.addComponent` fails.
## What the added tests cover
### 1. Pre-insertion failure for `MemoryTierWriter`
Test:
```scala
test("memory tier writer should preserve state on pre-insertion OOM")
```
This test verifies that if the failure happens before the buffer is actually
added to the flush buffer, then:
- `flushBuffer.writerIndex()` is unchanged
- `flushBuffer.numComponents()` is unchanged
- memory-file storage counter is unchanged
- disk-buffer counter is unchanged
- file length is unchanged
- the input `ByteBuf` reference count is unchanged
This is the “no ownership taken” case.
---
### 2. Composite buffer capacity overflow for `MemoryTierWriter`
Test:
```scala
test("memory tier writer should preserve state on composite buffer capacity
overflow")
```
This constructs a `CompositeByteBuf` that is close to `Int.MaxValue`
readable bytes and then tries to append more data. It verifies that when
`addComponent` fails because of capacity overflow:
- no bytes are accounted
- file length is unchanged
- the input buffer is not leaked
- the composite buffer state is unchanged
This is useful because capacity overflow is one realistic way `addComponent`
can fail before the new component becomes usable.
---
### 3. Post-insertion consolidation OOM for `MemoryTierWriter`
Test:
```scala
test("memory tier writer should account data added before consolidation OOM")
```
This is the most important regression case for the bug description.
It creates a `CompositeByteBuf` whose allocator fails when consolidation is
attempted. The test verifies that if the component has already been added
before consolidation fails, then:
- the flush buffer reflects the added bytes
- the memory-file storage counter is incremented
- the file length is incremented
- the input buffer has one additional reference owned by the flush buffer
- closing/releasing the memory file later releases the buffer and decrements
the counter correctly
This matches the intended behavior: if `addComponent` has already taken
ownership, the bytes should be accounted.
---
### 4. Post-insertion consolidation OOM for `LocalTierWriter`
Test:
```scala
test("local tier writer should account data added before consolidation OOM")
```
This is the local-disk equivalent of the previous test. It verifies that
when the component is added before consolidation fails:
- `MemoryManager.instance.incrementDiskBuffer(numBytes)` is reflected in the
disk-buffer counter
- returning the buffer later releases the reference and decrements the
counter
- the input buffer reference count returns to the expected value
This is important because `LocalTierWriter` and `DfsTierWriter` account
bytes through the disk-buffer counter, not the memory-file counter.
---
## Overall assessment of the tests
The added tests are good and directly target the regression described in the
PR:
- They distinguish between failure before insertion and failure after
insertion.
- They validate both storage accounting and reference-count behavior.
- They cover `MemoryTierWriter` and `LocalTierWriter`.
- They clean up global `MemoryManager` counters using `restoreCounters`,
which is important because those counters are global/static.
So the test direction is correct and the coverage is much better than having
only the happy-path tests.
---
## Suggested additional test coverage
There are a few gaps worth considering before merge.
### 1. Add a similar test for `DfsTierWriter`
The production change also modifies `DfsTierWriter`:
```scala
override protected def accountAddedBytes(numBytes: Int): Unit = {
MemoryManager.instance.incrementDiskBuffer(numBytes)
}
```
But the new tests only cover:
- `MemoryTierWriter`
- `LocalTierWriter`
Since `DfsTierWriter` has the same accounting path, it would be good to add:
```scala
test("dfs tier writer should account data added before consolidation OOM")
```
modeled after the local-tier test.
---
### 2. Add a true pre-insertion failure after `buf.retain()`
The current pre-insertion OOM test uses:
```scala
private class PreInsertionOomByteBuf(buffer: ByteBuf, failure:
OutOfMemoryError)
extends DuplicatedByteBuf(buffer) {
override def readerIndex(): Int = throw failure
}
```
However, in `appendToFlushBuffer`, the first operation is:
```scala
val numBytes = buf.readableBytes()
```
Because `readableBytes()` can call `readerIndex()`, this failure can happen
before:
```scala
buf.retain()
```
So the test does not fully exercise this path:
```scala
buf.retain()
flushBuffer.addComponent(true, buf) // fails before taking ownership
```
That path is important because the current code does:
```scala
buf.retain()
try {
flushBuffer.addComponent(true, buf)
} catch {
case t: Throwable if t.isInstanceOf[OutOfMemoryError] || NonFatal(t) =>
val componentAdded =
flushBuffer.writerIndex() != writerIndexBefore ||
flushBuffer.numComponents() != numComponentsBefore
if (componentAdded) {
accountAddedBytes(numBytes)
}
logError(...)
throw t
}
```
If `addComponent` fails after `buf.retain()` but before taking ownership,
the retained reference may not be released.
A stronger test would make `readableBytes()` succeed, but make the flush
buffer’s `addComponent` fail without changing `writerIndex` or `numComponents`.
For example, a custom `CompositeByteBuf` wrapper/subclass could throw from
`addComponent`.
Conceptually:
```scala
class FailingFlushBuffer extends CompositeByteBuf(
UnpooledByteBufAllocator.DEFAULT,
false,
16) {
override def addComponent(increaseWriterIndex: Boolean, buffer: ByteBuf):
CompositeByteBuf = {
throw new OutOfMemoryError("pre-insertion failure")
}
}
```
Then assert:
```scala
assert(buf.refCnt() === refCntBeforeWrite)
```
If that test fails, it would indicate that `appendToFlushBuffer` should
release the retained reference when `componentAdded == false`, for example:
```scala
if (componentAdded) {
accountAddedBytes(numBytes)
} else {
buf.release()
}
```
This is probably the most important follow-up test to consider.
---
### 3. Add a non-OOM `NonFatal` failure test
The catch block handles:
```scala
case t: Throwable if t.isInstanceOf[OutOfMemoryError] || NonFatal(t) =>
```
The tests currently focus on `OutOfMemoryError` and
`IllegalArgumentException`. It would also be useful to add a test where
`addComponent` throws a non-fatal runtime exception such as:
```scala
new IllegalStateException("synthetic failure")
```
This would verify that the same accounting/reference-count logic applies to
non-OOM failures.
---
### 4. Assert all captured counters in each test
Some tests capture both memory and disk counters but do not assert both.
For example, in the memory capacity-overflow test:
```scala
val diskCounterBefore = memoryManager.getDiskBufferCounter.get()
```
is captured, but the test later asserts only:
```scala
assert(memoryManager.getMemoryFileStorageCounter === memoryCounterBefore)
```
It would be slightly stronger to also assert:
```scala
assert(memoryManager.getDiskBufferCounter.get() === diskCounterBefore)
```
Similarly, the local-tier test captures:
```scala
val memoryCounterBefore = memoryManager.getMemoryFileStorageCounter
```
but does not assert that the memory counter remained unchanged. Adding that
assertion would make the tests more complete.
---
### 5. Consider reducing dependence on Netty exception messages
The capacity-overflow test asserts:
```scala
assert(thrown.getMessage.contains("overflow"))
```
This is reasonable, but it can become brittle if Netty changes the exception
message. If the important condition is the failure type, asserting only:
```scala
intercept[IllegalArgumentException]
```
may be enough. If the message check is kept, it should be treated as a
nice-to-have rather than the core validation.
---
## How to run the tests
The PR description says `TierWriterSuite` passes with 11 tests.
Depending on the repo’s Maven test configuration, the suite can typically be
run with something like:
```bash
./build/mvn test -pl worker
-DwildcardSuites=org.apache.celeborn.service.deploy.worker.storage.TierWriterSuite
```
or, if surefire/scalatest configuration requires disabling surefire:
```bash
./build/mvn test -pl worker -Dtest=none \
-DwildcardSuites=org.apache.celeborn.service.deploy.worker.storage.TierWriterSuite
```
And formatting can be checked with:
```bash
./build/mvn spotless:check -pl worker
```
---
## Summary
The tests added in this PR are relevant and target the actual bug: incorrect
accounting when `CompositeByteBuf.addComponent` fails. The post-insertion
consolidation OOM tests are especially valuable because they validate the
harder case where the buffer has already been accepted into the composite
buffer.
The main suggested improvements are:
1. Add a `DfsTierWriter` version of the consolidation-OOM test.
2. Add a test that forces failure after `buf.retain()` but before insertion.
3. Consider releasing the retained buffer when `componentAdded == false`.
4. Add a non-OOM `NonFatal` exception test.
5. Assert both memory and disk counters wherever both are captured.
--
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]
