laskoviymishka commented on code in PR #1368:
URL: https://github.com/apache/iceberg-go/pull/1368#discussion_r3535329291
##########
table/partitioned_fanout_writer_test.go:
##########
@@ -92,6 +93,101 @@ func (s *FanoutWriterTestSuite)
createCustomTestRecord(arrSchema *arrow.Schema,
return bldr.NewRecordBatch()
}
+func (s *FanoutWriterTestSuite) createLargeTestRecord(arrSchema *arrow.Schema,
rows int, idOffset int64, payloadSize int) arrow.RecordBatch {
+ bldr := array.NewRecordBuilder(s.mem, arrSchema)
+ defer bldr.Release()
+
+ payload := strings.Repeat("p", payloadSize)
+ for i := range rows {
+ bldr.Field(0).(*array.Int64Builder).Append(idOffset + int64(i))
+ bldr.Field(1).(*array.StringBuilder).Append(payload)
+ }
+
+ return bldr.NewRecordBatch()
+}
+
+func (s *FanoutWriterTestSuite)
TestCloseAllFlushesAfterFanoutSuccessContextCancel() {
Review Comment:
This test doesn't exercise the race it's guarding against. The name says
"ContextCancel" but nothing here cancels a context — it submits 10000 rows
through a single writer with a 256MB target and asserts they all land. On local
I/O the stream drains well before `fanoutWorkers.Wait()` returns, so the old
errgroup-ctx cancellation never catches an in-flight write. This passes on the
pre-fix code too, which means it isn't defending the fix.
To actually pin the bug I'd drop to the `RollingDataWriter` level: inject a
`FileWriter` (or a write-path hook) that blocks mid-write, cancel while the
stream is parked in it, then assert `closeAll()` flushes every record and
`abortAll()` discards them. That's the invariant this PR changes, and it's
doable entirely in `rolling_data_writer_test.go`. wdyt?
##########
table/partitioned_fanout_writer.go:
##########
@@ -115,7 +116,7 @@ func startRecordFeeder(ctx context.Context, itr
iter.Seq2[arrow.RecordBatch, err
})
}
-func (p *partitionedFanoutWriter) fanout(ctx context.Context, inputRecordsCh
<-chan arrow.RecordBatch, dataFilesChannel chan<- iceberg.DataFile) error {
+func (p *partitionedFanoutWriter) fanout(ctx context.Context, writerCtx
context.Context, inputRecordsCh <-chan arrow.RecordBatch, dataFilesChannel
chan<- iceberg.DataFile) error {
Review Comment:
Two adjacent `context.Context` params through
`fanout`/`processRecord`/`processBatch` makes me nervous — Go has no named args
at the call site, so a future reorder swaps them silently and the compiler
won't blink. And a swap here reintroduces exactly the race this PR is fixing.
Every test call already reads `processBatch(ctx, ctx, ...)`, which shows how
easy the confusion is.
Since `writerCtx` has the same lifetime as the `Write()` call, I'd attach it
to the `partitionedFanoutWriter` / `positionDeletePartitionedFanoutWriter`
struct during `Write()` and drop it from these signatures entirely. wdyt?
##########
table/rolling_data_writer.go:
##########
@@ -450,26 +451,64 @@ func (r *RollingDataWriter) sendError(err error) {
}
}
-func (r *RollingDataWriter) close() {
+// closeInput gracefully closes the record channel so any queued writes can
still be
+// processed while honoring context cancellation checks in the stream.
+func (r *RollingDataWriter) closeInput() {
+ r.closeRecordCh.Do(func() {
+ close(r.recordCh)
+ })
+}
+
+// abort cancels in-flight work and then closes input so the stream can exit.
+func (r *RollingDataWriter) abort() {
r.cancel()
- close(r.recordCh)
+ r.closeInput()
}
func (r *RollingDataWriter) closeAndWait() error {
- r.close()
+ r.closeInput()
r.factory.writers.Delete(r.partitionKey)
r.wg.Wait()
if err := <-r.errorCh; err != nil {
+ r.cancel()
+
return fmt.Errorf("error in rolling data writer: %w", err)
}
+ r.cancel()
+
return nil
}
+func (r *RollingDataWriter) abortAndWait() {
+ r.abort()
+ r.factory.writers.Delete(r.partitionKey)
+ r.wg.Wait()
+}
+
func (w *writerFactory) closeAll() error {
defer w.stopCount()
- var writers []*RollingDataWriter
+ writers := w.writerList()
+ var err error
+ for _, writer := range writers {
+ if closeErr := writer.closeAndWait(); closeErr != nil && err ==
nil {
+ err = closeErr
+ }
+ }
+
+ return err
+}
+
+func (w *writerFactory) abortAll() {
+ defer w.stopCount()
Review Comment:
Both `closeAll()` and `abortAll()` defer `w.stopCount()`, and the new tests
also `defer factory.closeAll()` on top of the normal path — so `stopCount()`
can run more than once per factory. It's harmless today because `iter.Pull`'s
stop is idempotent, but the ownership is implicit and would turn into a real
bug the moment `stopCount` wraps anything non-idempotent. Might be worth
pulling the `stopCount()` responsibility to a single owner.
##########
table/rolling_data_writer.go:
##########
@@ -450,26 +451,64 @@ func (r *RollingDataWriter) sendError(err error) {
}
}
-func (r *RollingDataWriter) close() {
+// closeInput gracefully closes the record channel so any queued writes can
still be
+// processed while honoring context cancellation checks in the stream.
+func (r *RollingDataWriter) closeInput() {
+ r.closeRecordCh.Do(func() {
+ close(r.recordCh)
+ })
+}
+
+// abort cancels in-flight work and then closes input so the stream can exit.
+func (r *RollingDataWriter) abort() {
r.cancel()
- close(r.recordCh)
+ r.closeInput()
}
func (r *RollingDataWriter) closeAndWait() error {
- r.close()
+ r.closeInput()
r.factory.writers.Delete(r.partitionKey)
r.wg.Wait()
if err := <-r.errorCh; err != nil {
+ r.cancel()
+
return fmt.Errorf("error in rolling data writer: %w", err)
}
+ r.cancel()
+
return nil
}
+func (r *RollingDataWriter) abortAndWait() {
+ r.abort()
+ r.factory.writers.Delete(r.partitionKey)
+ r.wg.Wait()
Review Comment:
`abortAndWait()` never drains `r.errorCh` after `Wait()`, so if `stream()`
sent an error it's silently dropped. No leak since `errorCh` is buffered at 1,
but it makes a failed write harder to debug. Even just draining and logging it
would help.
##########
table/rolling_data_writer.go:
##########
@@ -479,12 +518,5 @@ func (w *writerFactory) closeAll() error {
return true
})
- var err error
- for _, writer := range writers {
- if closeErr := writer.closeAndWait(); closeErr != nil && err ==
nil {
- err = closeErr
- }
- }
-
- return err
+ return
Review Comment:
nlreturn is enabled in `.golangci.yml`, and a naked `return` on a named
result tends to trip it — worth confirming this doesn't fail CI. Either way I'd
return the local explicitly (`return writers`) or drop the named return; it
reads better than the bare form here.
##########
table/rolling_data_writer.go:
##########
@@ -450,26 +451,64 @@ func (r *RollingDataWriter) sendError(err error) {
}
}
-func (r *RollingDataWriter) close() {
+// closeInput gracefully closes the record channel so any queued writes can
still be
+// processed while honoring context cancellation checks in the stream.
+func (r *RollingDataWriter) closeInput() {
+ r.closeRecordCh.Do(func() {
+ close(r.recordCh)
+ })
+}
+
+// abort cancels in-flight work and then closes input so the stream can exit.
+func (r *RollingDataWriter) abort() {
r.cancel()
- close(r.recordCh)
+ r.closeInput()
}
func (r *RollingDataWriter) closeAndWait() error {
- r.close()
+ r.closeInput()
r.factory.writers.Delete(r.partitionKey)
r.wg.Wait()
if err := <-r.errorCh; err != nil {
+ r.cancel()
+
return fmt.Errorf("error in rolling data writer: %w", err)
}
+ r.cancel()
Review Comment:
Both branches call `r.cancel()` right before returning. `CancelFunc` is
idempotent so it's safe, but a single `defer r.cancel()` at the top of
`closeAndWait()` reads cleaner and won't get dropped by a future early return.
##########
table/rolling_data_writer_test.go:
##########
@@ -234,6 +234,67 @@ func (s *RollingDataWriterTestSuite)
TestNewWriterFactoryReturnsErrorForInvalidF
s.ErrorContains(err, "withFactoryFileSchema")
}
+func (s *RollingDataWriterTestSuite)
TestCloseAllFinishesQueuedRecordsWithoutCancellingContext() {
+ 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)
+
+ outputCh := make(chan iceberg.DataFile, 10)
+ writerCtx, cancel := context.WithCancel(s.ctx)
+ defer cancel()
+
+ writer, err := factory.getOrCreateRollingDataWriter(writerCtx, "", nil,
outputCh)
+ s.Require().NoError(err)
+ const totalBatches = 4
+ const rowsPerBatch = 25
+ var expectedRows int64
+ for range totalBatches {
+ record := s.buildRecord(arrSchema, rowsPerBatch)
+ expectedRows += record.NumRows()
+ s.Require().NoError(writer.Add(record))
+ record.Release()
+ }
+
+ s.Require().NoError(factory.closeAll())
+ s.Require().Nil(writerCtx.Err(), "normal close should finish queued
writes without aborting context")
+
+ close(outputCh)
+ var actualRows int64
+ for df := range outputCh {
+ actualRows += df.Count()
+ }
+
+ s.Equal(expectedRows, actualRows, "all queued records should be flushed
when closeAll is used")
+}
+
+func (s *RollingDataWriterTestSuite) TestAbortAndWaitCancelsContext() {
+ 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)
+ defer factory.closeAll()
+
+ outputCh := make(chan iceberg.DataFile, 10)
+ writerCtx, cancel := context.WithCancel(s.ctx)
+ defer cancel()
+
+ writer, err := factory.getOrCreateRollingDataWriter(writerCtx, "", nil,
outputCh)
+ s.Require().NoError(err)
+ record := s.buildRecord(arrSchema, 5)
+ s.Require().NoError(writer.Add(record))
+ record.Release()
+
+ writer.abortAndWait()
+ s.Require().ErrorIs(writer.ctx.Err(), context.Canceled)
+}
Review Comment:
This only checks the context flipped to cancelled — it reaches into the
unexported `ctx` field and never verifies the incomplete Parquet file was
actually cleaned up. Walking the temp dir afterward (like
`TestStreamErrorPathUsesAbort` does) would prove abort went through `Abort()`
rather than `Close()`, which is the behavior that matters. Asserting on a
public-ish signal instead of `writer.ctx` would also loosen the white-box
coupling.
##########
table/partitioned_fanout_writer.go:
##########
@@ -174,18 +175,37 @@ func (p *partitionedFanoutWriter) processRecord(ctx
context.Context, record arro
return nil
}
-func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile)
iter.Seq2[iceberg.DataFile, error] {
- return yieldDataFiles(p.writerFactory, fanoutWorkers, outputDataFilesCh)
+func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile, writerCancel
context.CancelFunc) iter.Seq2[iceberg.DataFile, error] {
+ return yieldDataFiles(
+ p.writerFactory,
+ fanoutWorkers,
+ outputDataFilesCh,
+ p.writerFactory.closeAll,
+ p.writerFactory.abortAll,
+ writerCancel,
+ )
}
-func yieldDataFiles(writerFactory *writerFactory, fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile)
iter.Seq2[iceberg.DataFile, error] {
+func yieldDataFiles(
Review Comment:
`writerFactory` is dead now that `closeAll`/`abortAll` come in as explicit
callbacks — the body doesn't touch it anymore. Both call sites already pass the
two closures, so I'd just drop the param.
##########
table/partitioned_fanout_writer.go:
##########
@@ -174,18 +175,37 @@ func (p *partitionedFanoutWriter) processRecord(ctx
context.Context, record arro
return nil
}
-func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile)
iter.Seq2[iceberg.DataFile, error] {
- return yieldDataFiles(p.writerFactory, fanoutWorkers, outputDataFilesCh)
+func (p *partitionedFanoutWriter) yieldDataFiles(fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile, writerCancel
context.CancelFunc) iter.Seq2[iceberg.DataFile, error] {
+ return yieldDataFiles(
+ p.writerFactory,
+ fanoutWorkers,
+ outputDataFilesCh,
+ p.writerFactory.closeAll,
+ p.writerFactory.abortAll,
+ writerCancel,
+ )
}
-func yieldDataFiles(writerFactory *writerFactory, fanoutWorkers
*errgroup.Group, outputDataFilesCh chan iceberg.DataFile)
iter.Seq2[iceberg.DataFile, error] {
+func yieldDataFiles(
+ writerFactory *writerFactory,
+ fanoutWorkers *errgroup.Group,
+ outputDataFilesCh chan iceberg.DataFile,
+ closeAll func() error,
+ abortAll func(),
+ writerCancel context.CancelFunc,
+) iter.Seq2[iceberg.DataFile, error] {
// Use a channel to safely communicate the error from the goroutine
// to avoid a data race between writing err in the goroutine and
reading it in the iterator.
errCh := make(chan error, 1)
go func() {
defer close(outputDataFilesCh)
Review Comment:
The LIFO ordering here is load-bearing — `writerCancel()` has to run before
`close(outputDataFilesCh)`, and reversing these two lines silently reverses the
semantics with nothing to catch it. A one-line comment on why the order matters
would save the next reader.
--
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]
