laskoviymishka commented on code in PR #1851:
URL: https://github.com/apache/iceberg-go/pull/1851#discussion_r3833959796
##########
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:
This defer runs on every `stream()` exit, including the clean channel-closed
path, so `r.cancel()` now fires from the stream goroutine on success too,
before `wg.Done()` unblocks `closeAndWait`. It's idempotent today
(`closeAndWait` cancels again), but it's a quiet semantic shift: the per-writer
context is now cancelled "by stream" on the happy path, not "by closeAndWait",
so anything that later checks `r.ctx.Err()` after `wg.Wait()` returns would see
a cancelled context on a successful write. If we keep this shape I'd guard it
behind an exited-cleanly flag so the cancel only fires on the error exit.
##########
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
Review Comment:
Small thing, but "recordCh is never closed here" reads like the `!ok` case a
few lines down is dead code. It isn't: `closeInput()` closes the channel and,
since `wg.Done()` is the last defer here (LIFO), it can run while this drain is
still looping. So `!ok` is genuinely reachable, and someone could prune it as
dead code and turn the drain into a spin on a closed channel. I'd reword to say
another goroutine may close `recordCh` via `closeInput`/abort and that `!ok` is
what ends the drain.
While here: this comment block is much denser than the rest of the file,
where the method bodies are mostly one line or none. I'd trim it to the intent
plus the one real constraint (don't close recordCh here) and let the code carry
the rest.
##########
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)
Review Comment:
Missing the `defer cancel()` that both other `WithCancel` tests in this file
pair with. `stream`'s cleanup does call `r.cancel()`, but only once the
goroutine is running, so if the test bails between here and `go
writer.stream(...)` the context leaks (and `go vet`/lint flags the missing
cancel). One line right below.
##########
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)
Review Comment:
`createWriterFactory` calls `iter.Pull`, which starts a goroutine that only
stops via `stopCount()` inside `closeAll()`/`abortAll()`. Every other test in
this file that builds a factory pairs it with `defer factory.closeAll()`; this
one doesn't, so the pull goroutine leaks for the rest of the run (a goleak
check would flag it). I'd add `defer factory.closeAll()` right here. No writer
is registered on this factory, so it just calls `stopCount()`.
--
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]
