diff --git a/out_writeapi.go b/out_writeapi.go index caeec5d..b61427d 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -65,6 +65,13 @@ type outputConfig struct { timestampFields []string fieldCache fieldLookupCache flushTimeout time.Duration + // instanceCtx is a long-lived context, scoped to the lifetime of the output instance, + // used for AppendRows and managed stream creation. Unlike the per-flush context, it is not + // canceled at the end of a flush, so the managed writer's internal retry/reconnect logic for + // asynchronously-acknowledged appends (default/at-least-once path) is not aborted prematurely. + // It is canceled in FLBPluginExitCtx after pending responses are drained and streams are closed. + instanceCtx context.Context + instanceCancel context.CancelFunc } var ( @@ -424,7 +431,7 @@ func sendRequestRetries(ctx context.Context, data [][]byte, config *outputConfig currStream.managedstream.Finalize(ctx) currStream.managedstream.Close() // Rebuild stream - buildErr := buildStream(ctx, config, streamIndex) + buildErr := buildStream(config.instanceCtx, config, streamIndex) config.streamMu.Unlock() if buildErr != nil { return buildErr @@ -444,7 +451,13 @@ func sendRequestRetries(ctx context.Context, data [][]byte, config *outputConfig return nil } -// This function sends data and appends the responses to a queue to be checked asynchronously through a default stream with at least once functionality +// This function sends data and appends the responses to a queue to be checked asynchronously through a default stream with at least once functionality. +// The ctx parameter is retained for signature symmetry with sendRequestRetries (both are dispatched +// from sendRequest); it is intentionally not used for the append itself. The append uses the +// instance-scoped context instead. Note that because instanceCtx has no per-append deadline, the +// only backpressure bound is the managed writer's inflight limit (Max_Queue_Requests / +// Max_Queue_Bytes); on saturation AppendRows blocks rather than timing out, which surfaces as a +// blocking flush absorbed by Fluent Bit's (filesystem) buffer. func sendRequestDefault(ctx context.Context, data [][]byte, config *outputConfig, streamIndex int) error { // Hold streamMu only for the slice access; AppendRows is goroutine-safe inside // the managed writer and must not be called under the lock to avoid blocking all @@ -452,7 +465,10 @@ func sendRequestDefault(ctx context.Context, data [][]byte, config *outputConfig // is saturated. currStream := config.getStream(streamIndex) - appendResult, err := currStream.managedstream.AppendRows(ctx, data) + // Use the instance-scoped context (not the per-flush ctx) so that the managed writer's + // internal retry/reconnect for this asynchronously-acknowledged append is not aborted + // when the originating flush's context is canceled at flush end. + appendResult, err := currStream.managedstream.AppendRows(config.instanceCtx, data) if err != nil { return err } @@ -536,8 +552,10 @@ var setThreshold = func(maxQueueSize int) int { } // This function check whether there is room for scaling and the scales the number of stream dynamically depending on if it -// Detects back pressure from the queue -func createNewStreamDynamicScaling(ctx context.Context, config *outputConfig) { +// Detects back pressure from the queue. +// It uses config.instanceCtx (not a per-flush context) for stream creation so that the new stream's +// lifetime is not tied to the flush that triggered the scaling. +func createNewStreamDynamicScaling(config *outputConfig) { config.streamMu.Lock() defer config.streamMu.Unlock() if len(*config.managedStreamSlice) < maxNumStreamsPerInstance { @@ -554,7 +572,10 @@ func createNewStreamDynamicScaling(ctx context.Context, config *outputConfig) { if mostEfficientQueueLength > config.requestCountThreshold { *config.managedStreamSlice = append(*config.managedStreamSlice, &newStream) newStreamIndex := len(*config.managedStreamSlice) - 1 - err := buildStream(ctx, config, newStreamIndex) + // Use the instance-scoped context so the newly created stream's lifetime is not tied + // to the per-flush context (which is canceled at flush end and would otherwise + // terminalize the new managed stream immediately). + err := buildStream(config.instanceCtx, config, newStreamIndex) if err != nil { log.Printf("Creating an additional managed stream with destination table: %s failed in FLBPluginInit: %s", config.tableRef, err) // If failure, failed stream is removed from slice @@ -731,9 +752,13 @@ func FLBPluginInit(plugin unsafe.Pointer) int { } // Create new client - initCtx := context.Background() - client, err := getClient(initCtx, projectID) + // Create an instance-scoped context that outlives individual flush calls. Using it for client + // creation, stream creation (buildStream) and AppendRows ensures the managed writer's internal + // retry/reconnect logic is not aborted when a per-flush context is canceled at flush end. + instanceCtx, instanceCancel := context.WithCancel(context.Background()) + client, err := getClient(instanceCtx, projectID) if err != nil { + instanceCancel() log.Printf("Creating a new managed BigQuery Storage write client scoped to: %s failed in FLBPluginInit: %s", projectID, err) return output.FLB_ERROR } @@ -742,8 +767,9 @@ func FLBPluginInit(plugin unsafe.Pointer) int { tableReference := fmt.Sprintf("projects/%s/datasets/%s/tables/%s", projectID, datasetID, tableID) // Call getDescriptors to get the message descriptor, and descriptor proto - md, descriptor, timestampFields, err := getDescriptors(initCtx, client, projectID, datasetID, tableID, dateTimeStringType) + md, descriptor, timestampFields, err := getDescriptors(instanceCtx, client, projectID, datasetID, tableID, dateTimeStringType) if err != nil { + instanceCancel() log.Printf("Getting message descriptor and descriptor proto for table: %s failed in FLBPluginInit: %s", tableReference, err) return output.FLB_ERROR } @@ -788,11 +814,14 @@ func FLBPluginInit(plugin unsafe.Pointer) int { timestampFields: timestampFields, fieldCache: buildFieldLookupCache(md), flushTimeout: time.Duration(flushTimeoutSec) * time.Second, + instanceCtx: instanceCtx, + instanceCancel: instanceCancel, } // Create stream using NewManagedStream - err = buildStream(initCtx, &config, 0) + err = buildStream(instanceCtx, &config, 0) if err != nil { + instanceCancel() log.Printf("Creating a new managed stream with destination table: %s failed in FLBPluginInit: %s", tableReference, err) return output.FLB_ERROR } @@ -877,7 +906,7 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int // Drain ready responses and check whether a new stream should be created. checkAllStreamResponses(flushCtx, &config.managedStreamSlice, false, &config.streamMu, config.exactlyOnce, id) - createNewStreamDynamicScaling(flushCtx, config) + createNewStreamDynamicScaling(config) // binaryData is a flush-local buffer. Each worker (workers > 1) gets its own // stack frame, so there is no shared state and no mutex needed here. @@ -925,6 +954,13 @@ func FLBPluginExitCtx(ctx unsafe.Pointer) int { } } + // Cancel the instance-scoped context only after pending responses are drained and the streams + // and client are closed, so that in-flight AppendRows are not aborted prematurely (which would + // cause data loss and terminalize the streams during shutdown). + if config.instanceCancel != nil { + config.instanceCancel() + } + if errFlag { return output.FLB_ERROR } diff --git a/out_writeapi_test.go b/out_writeapi_test.go index 683f756..d92d169 100644 --- a/out_writeapi_test.go +++ b/out_writeapi_test.go @@ -1303,3 +1303,123 @@ func TestConvertTimestampFieldsRaw(t *testing.T) { }) } } + +// TestFLBPluginFlushCtxAppendUsesInstanceContext is a regression test for the bug where the +// per-flush context was passed to AppendRows on the default (at-least-once) path. Because appends +// are acknowledged asynchronously on a later flush, canceling the flush context at flush end aborted +// the managed writer's internal retry/reconnect, turning transient connection failures into a +// permanent failure storm. This test verifies that: +// - AppendRows receives the instance-scoped context (config.instanceCtx), not the flush context. +// - That context is NOT canceled after FLBPluginFlushCtx returns (when the flush context is canceled). +// - The instance-scoped context IS canceled after FLBPluginExitCtx. +func TestFLBPluginFlushCtxAppendUsesInstanceContext(t *testing.T) { + var setID int + + testTableSchema := &storagepb.TableSchema{ + Fields: []*storagepb.TableFieldSchema{ + {Name: "Text", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE}, + }, + } + + mockClient := &MockManagedWriterClient{ + NewManagedStreamFunc: func(ctx context.Context, opts ...managedwriter.WriterOption) (*managedwriter.ManagedStream, error) { + return nil, nil + }, + GetWriteStreamFunc: func(ctx context.Context, req *storagepb.GetWriteStreamRequest, opts ...gax.CallOption) (*storagepb.WriteStream, error) { + return &storagepb.WriteStream{Name: "mockstream", TableSchema: testTableSchema}, nil + }, + CloseFunc: func() error { return nil }, + } + + originalGetClient := getClient + getClient = func(ctx context.Context, projectID string) (ManagedWriterClient, error) { + return mockClient, nil + } + defer func() { getClient = originalGetClient }() + + // Capture the context that AppendRows is invoked with. + var capturedCtx context.Context + mockMS := &MockManagedStream{ + AppendRowsFunc: func(ctx context.Context, data [][]byte, opts ...managedwriter.AppendOption) (*managedwriter.AppendResult, error) { + capturedCtx = ctx + return nil, nil // nil AppendResult queued; drained via mocked pluginGetResult at exit + }, + CloseFunc: func() error { return nil }, + FinalizeFunc: func(ctx context.Context, opts ...gax.CallOption) (int64, error) { + return 0, nil + }, + FlushRowsFunc: func(ctx context.Context, offset int64, opts ...gax.CallOption) (int64, error) { + return 0, nil + }, + StreamNameFunc: func() string { return "" }, + } + + origGetWriter := getWriter + getWriter = func(client ManagedWriterClient, ctx context.Context, projectID string, opts ...managedwriter.WriterOption) (MWManagedStream, error) { + return mockMS, nil + } + defer func() { getWriter = origGetWriter }() + + patch1 := gomonkey.ApplyFunc(output.FLBPluginConfigKey, func(plugin unsafe.Pointer, key string) string { + return "" + }) + defer patch1.Reset() + + patchSetContext := gomonkey.ApplyFunc(output.FLBPluginSetContext, func(plugin unsafe.Pointer, ctx interface{}) { + setID = ctx.(int) + }) + defer patchSetContext.Reset() + + initRes := FLBPluginInit(nil) + assert.Equal(t, output.FLB_OK, initRes) + + origGetContext := getFLBPluginContext + getFLBPluginContext = func(ctx unsafe.Pointer) int { + if ctx != nil { + return *(*int)(ctx) + } + return 0 + } + defer func() { getFLBPluginContext = origGetContext }() + + // Keep queued responses "in-flight" so the flush does not try to drain them. + origIsReady := isReady + isReady = func(_ *managedwriter.AppendResult) bool { return false } + defer func() { isReady = origIsReady }() + + // Avoid dereferencing the nil AppendResult during the exit drain. + origGetResult := pluginGetResult + pluginGetResult = func(_ *managedwriter.AppendResult, _ context.Context) (int64, error) { + return -1, nil + } + defer func() { pluginGetResult = origGetResult }() + + origDecoder := newDecoder + newDecoder = func(data unsafe.Pointer, length int) *output.FLBDecoder { return nil } + defer func() { newDecoder = origDecoder }() + + var rowSent bool + patchRecord := gomonkey.ApplyFunc(output.GetRecord, func(dec *output.FLBDecoder) (int, interface{}, map[interface{}]interface{}) { + if !rowSent { + rowSent = true + return 0, nil, map[interface{}]interface{}{"Text": []byte("hello")} + } + return 1, nil, nil + }) + defer patchRecord.Reset() + + pointerValue := unsafe.Pointer(&setID) + flushResult := FLBPluginFlushCtx(pointerValue, nil, 0, nil) + assert.Equal(t, output.FLB_OK, flushResult) + + // The per-flush context is canceled when FLBPluginFlushCtx returns. The append must have + // received the instance-scoped context, which stays alive until FLBPluginExitCtx. + assert.NotNil(t, capturedCtx, "AppendRows should have been called") + assert.NoError(t, capturedCtx.Err(), "AppendRows context must not be canceled after the flush returns") + assert.Equal(t, configMap[setID].instanceCtx, capturedCtx, "AppendRows must use the instance-scoped context") + + // After exit, the instance-scoped context is canceled. + exitResult := FLBPluginExitCtx(pointerValue) + assert.Equal(t, output.FLB_OK, exitResult) + assert.Error(t, capturedCtx.Err(), "instance-scoped context must be canceled after exit") +}