zeroshade commented on code in PR #1851:
URL: https://github.com/apache/iceberg-go/pull/1851#discussion_r3815123078
##########
table/rolling_data_writer.go:
##########
@@ -430,6 +442,31 @@ func (r *RollingDataWriter) stream(outputDataFilesCh
chan<- iceberg.DataFile) {
if currentWriter != nil {
_ = currentWriter.Abort()
}
+ // stream is exiting; this defer runs before the
CompareAndDelete
+ // defer above deregisters the writer, so cleanup does not
depend on
+ // abortAll finding this writer in the registry: by the time (if
+ // ever) abortAll looks, this writer is already gone. Cancel
first so
+ // any Add call still in flight from another fanout worker
prefers
+ // bailing out via ctx.Done() (see Add's priority check) over
queuing
+ // into a channel nobody will drain again, then release
whatever is
+ // already buffered right now. recordCh is never closed here:
other
+ // callers may still be sending to it, and closing while a send
can
+ // still land would risk a send-on-closed-channel panic. That
leaves
+ // a narrow window between this drain and Add's ctx.Done() check
+ // taking effect elsewhere, but it is bounded to that one window
+ // rather than persisting for the life of the writer.
+ r.cancel()
Review Comment:
**Blocking: an in-flight `Add` can enqueue after this drain finishes.** The
first `ctx.Done()` check does not synchronize with cleanup. If `Add` passes
that check and pauses in `record.Retain()`, this defer can cancel, observe an
empty `recordCh`, and return. When `Add` resumes, both the buffered send and
`ctx.Done()` are ready, so Go may choose the send; `Add` returns `nil` and the
retained record remains queued with no consumer. I reproduced this controlled
interleaving 200 times and observed 108 post-drain sends. The acknowledged
window therefore still preserves the leak this PR is intended to fix. Please
add a barrier with active `Add` calls (or otherwise make send-vs-shutdown
atomic) before the final drain completes.
##########
table/rolling_data_writer_test.go:
##########
@@ -905,3 +905,60 @@ func (s *RollingDataWriterTestSuite)
TestStreamRecoversWriterClosePanic() {
})
}
}
+
+// failingOpenFormat always fails to open a file writer, simulating a
+// mid-stream write error (e.g. an unwritable location) that makes stream
+// return before it ever dequeues most of what's already buffered in recordCh.
+type failingOpenFormat struct {
+ tblutils.FileFormat
+}
+
+func (failingOpenFormat) NewFileWriter(context.Context, iceio.WriteFileIO,
map[int]any, tblutils.WriteFileInfo, *arrow.Schema) (tblutils.FileWriter,
error) {
+ return nil, errors.New("simulated open failure")
+}
+
+// TestStreamErrorDrainsBufferedRecords reproduces the leak from #1825: when
+// stream exits early on a write error, records already sitting in recordCh —
+// queued by Add before stream's first dequeue attempt fails — must still be
+// released. Before the fix, stream's deferred CompareAndDelete deregisters the
+// writer before abortAll ever has a chance to run, so nothing drains what's
+// left in recordCh. Records are queued directly, bypassing Add and the
+// goroutine start in newRollingDataWriter, so all of them are buffered before
+// stream ever runs — otherwise Add would race stream's error exit for the
+// later records once errorCh closes.
+func (s *RollingDataWriterTestSuite) TestStreamErrorDrainsBufferedRecords() {
+ arrSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+
+ loc := filepath.ToSlash(s.T().TempDir())
+ factory, _ := s.createWriterFactory(loc, arrSchema, 1024*1024)
+ factory.format = failingOpenFormat{FileFormat: factory.format}
+
+ outputCh := make(chan iceberg.DataFile, 10)
+ ctx, cancel := context.WithCancel(s.ctx)
+ writer := &RollingDataWriter{
+ recordCh: make(chan arrow.RecordBatch, 64),
+ errorCh: make(chan error, 1),
+ factory: factory,
+ ctx: ctx,
+ cancel: cancel,
+ }
+
+ const numRecords = 5
+ for range numRecords {
+ record := s.buildRecord(arrSchema, 3)
+ record.Retain()
+ writer.recordCh <- record
Review Comment:
This setup proves that cleanup releases records already buffered before the
stream exits, but bypassing `Add` also bypasses the failing interleaving.
Please exercise an `Add` that is in flight between its initial context check
and its send while the stream shuts down, and assert that it cannot return
success or leave a record in `recordCh` after cleanup.
--
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]
