From 57c7aa811aa742595e7d85ea11fcc7deb15bbeeb Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 17:55:34 +0900 Subject: [PATCH 01/12] fix: support nested TIMESTAMP field conversion in RECORD types --- out_writeapi.go | 20 ++++++++++++++++---- out_writeapi_test.go | 24 ++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/out_writeapi.go b/out_writeapi.go index 9c8474b..d7ff644 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -21,6 +21,7 @@ import ( "log" "math" "strconv" + "strings" "sync" "time" "unsafe" @@ -187,12 +188,23 @@ func getDescriptors(curr_ctx context.Context, mw_client ManagedWriterClient, pro // convertTimestampFieldsRaw converts timestamp fields in a raw Fluent Bit record // (map[interface{}]interface{}) from seconds to microseconds for BigQuery Storage Write API. // BigQuery TIMESTAMP type expects microseconds since Unix epoch. +// Field paths may be dotted (e.g. "outer.inner") for nested RECORD fields. func convertTimestampFieldsRaw(data map[interface{}]interface{}, timestampFields []string) { for _, field := range timestampFields { - // In Go, map[interface{}]interface{} lookup with a string key works - // because interface comparison uses underlying type+value equality. - if val, ok := data[field]; ok { - convertTimestampValueRaw(data, field, val) + if idx := strings.IndexByte(field, '.'); idx >= 0 { + // Nested path: navigate into the sub-map and recurse with the remainder. + parent, remainder := field[:idx], field[idx+1:] + if sub, ok := data[parent]; ok { + if subMap, ok := sub.(map[interface{}]interface{}); ok { + convertTimestampFieldsRaw(subMap, []string{remainder}) + } + } + } else { + // In Go, map[interface{}]interface{} lookup with a string key works + // because interface comparison uses underlying type+value equality. + if val, ok := data[field]; ok { + convertTimestampValueRaw(data, field, val) + } } } } diff --git a/out_writeapi_test.go b/out_writeapi_test.go index f71ca5e..8102629 100644 --- a/out_writeapi_test.go +++ b/out_writeapi_test.go @@ -1115,6 +1115,30 @@ func TestConvertTimestampFieldsRaw(t *testing.T) { timestampFields: []string{"time", "created_at"}, expected: map[interface{}]interface{}{"time": int64(1700000000000000), "created_at": int64(1600000000000000)}, }, + { + name: "nested RECORD timestamp field converted", + input: map[interface{}]interface{}{ + "meta": map[interface{}]interface{}{ + "ts": int64(1700000000), + }, + }, + timestampFields: []string{"meta.ts"}, + expected: map[interface{}]interface{}{ + "meta": map[interface{}]interface{}{ + "ts": int64(1700000000000000), + }, + }, + }, + { + name: "nested RECORD timestamp field missing parent not panic", + input: map[interface{}]interface{}{ + "other": int64(1700000000), + }, + timestampFields: []string{"meta.ts"}, + expected: map[interface{}]interface{}{ + "other": int64(1700000000), + }, + }, { name: "field not present in data", input: map[interface{}]interface{}{"other": "value"}, From 986b412f53ae478bbe35577ce549461555c4106a Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 18:15:22 +0900 Subject: [PATCH 02/12] fix: re-evaluate least loaded stream per chunk in FLBPluginFlushCtx --- out_writeapi.go | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/out_writeapi.go b/out_writeapi.go index d7ff644..a7d9b86 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -459,11 +459,11 @@ func getInstanceCount() int { // Finds the stream index when dynamically scaling func getLeastLoadedStream(streamSlice *[]*streamConfig) int { - min := len(*(*streamSlice)[0].appendResults) + m := len(*(*streamSlice)[0].appendResults) minStreamIndex := 0 for streamIndex, stream := range *streamSlice { - if len(*stream.appendResults) < min { - min = len(*stream.appendResults) + if len(*stream.appendResults) < m { + m = len(*stream.appendResults) minStreamIndex = streamIndex } } @@ -771,11 +771,6 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int // Keeps track of the number of rows previously sent var rowCounter int64 - // Find stream with least number of awaiting queue responses - config.mutex.Lock() - leastLoadedStreamIndex := getLeastLoadedStream(config.managedStreamSlice) - config.mutex.Unlock() - // Iterate Records for { // Extract Record @@ -794,6 +789,11 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int } else { // Successful data transformation if (currsize + len(buf)) >= config.maxChunkSize { + // Re-evaluate the least loaded stream per chunk to distribute load across streams. + config.mutex.Lock() + leastLoadedStreamIndex := getLeastLoadedStream(config.managedStreamSlice) + config.mutex.Unlock() + // Appending Rows err := sendRequest(ms_ctx, binaryData, &config, leastLoadedStreamIndex) if err != nil { @@ -817,6 +817,11 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int rowCounter++ } } + // Re-evaluate the least loaded stream for the final chunk. + config.mutex.Lock() + leastLoadedStreamIndex := getLeastLoadedStream(config.managedStreamSlice) + config.mutex.Unlock() + // Appending Rows err := sendRequest(ms_ctx, binaryData, &config, leastLoadedStreamIndex) if err != nil { From cfac0ae065bff8a2c1accade51ddcc0702efa6b6 Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 18:15:45 +0900 Subject: [PATCH 03/12] fix test code --- map_to_proto_test.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/map_to_proto_test.go b/map_to_proto_test.go index 583544d..4d5daca 100644 --- a/map_to_proto_test.go +++ b/map_to_proto_test.go @@ -880,7 +880,14 @@ func TestRawMapToBinary_MatchesMapToBinary(t *testing.T) { parsedBytes, err := mapToBinary(md, parsedData, cache) require.NoError(t, err) - assert.Equal(t, parsedBytes, rawBytes) + // Compare decoded messages instead of raw bytes: map iteration order is + // non-deterministic, so the two serializations may encode fields in + // different orders while remaining logically identical. + rawMsg := dynamicpb.NewMessage(md) + require.NoError(t, proto.Unmarshal(rawBytes, rawMsg)) + parsedMsg := dynamicpb.NewMessage(md) + require.NoError(t, proto.Unmarshal(parsedBytes, parsedMsg)) + assert.True(t, proto.Equal(rawMsg, parsedMsg)) } // TestRawMapToBinary_BoolWithBytes tests []byte to bool conversion. From 73f6efda5276f1091b8d8d8013b973f42dee28e5 Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 18:28:55 +0900 Subject: [PATCH 04/12] fix: replace gomonkey patch on output.NewDecoder with injectable function variable gomonkey's ApplyFunc for output.NewDecoder was unreliable on ARM64 macOS when tests ran sequentially in the same process. The patch occasionally failed to intercept calls, causing checks.createDecoder to stay at 0 and TestFLBPluginFlushCtxErrorHandling to fail non-deterministically. Fix: introduce a var newDecoder function variable in out_writeapi.go (matching the pattern already used for getClient, getWriter, getFLBPluginContext, pluginGetResult, isReady). Replace all four gomonkey.ApplyFunc(output.NewDecoder, ...) calls in tests with direct assignment to newDecoder, with deferred restore. This eliminates the ARM64 code-patching dependency entirely. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- out_writeapi.go | 32 +++++++++++++++++++++++++------- out_writeapi_test.go | 28 ++++++++++++++++------------ 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/out_writeapi.go b/out_writeapi.go index a7d9b86..c037ed1 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -63,6 +63,7 @@ type outputConfig struct { numRetries int timestampFields []string fieldCache fieldLookupCache + flushTimeout time.Duration } var ( @@ -82,6 +83,7 @@ const ( minQueueRequests = 10 dateTimeDefault = true maxUnixSeconds = 4102444800 // 2100-01-01 00:00:00 UTC in seconds + flushTimeoutSecDefault = 60 ) // This function mangles the top-level and complex (struct) BigQuery schema to convert NUMERIC, BIGNUMERIC, DATETIME, TIME, and JSON fields to STRING. @@ -488,7 +490,7 @@ 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(config **outputConfig) { +func createNewStreamDynamicScaling(ctx context.Context, config **outputConfig) { (*config).mutex.Lock() defer (*config).mutex.Unlock() if len(*(*config).managedStreamSlice) < maxNumStreamsPerInstance { @@ -505,7 +507,7 @@ func createNewStreamDynamicScaling(config **outputConfig) { if mostEfficientQueueLength > (*config).requestCountThreshold { *(*config).managedStreamSlice = append(*(*config).managedStreamSlice, &newStream) newStreamIndex := len(*(*config).managedStreamSlice) - 1 - err := buildStream(ms_ctx, config, newStreamIndex) + err := buildStream(ctx, 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 @@ -574,6 +576,11 @@ var getWriter = func(client ManagedWriterClient, ctx context.Context, projectID return client.NewManagedStream(ctx, opts...) } +// newDecoder is a wrapper around output.NewDecoder to allow test injection. +var newDecoder = func(data unsafe.Pointer, length int) *output.FLBDecoder { + return output.NewDecoder(data, length) +} + // This function acts as a wrapper for the GetContext function so that we may override it to // Mock it whenever needed var getFLBPluginContext = func(ctx unsafe.Pointer) int { @@ -664,6 +671,13 @@ func FLBPluginInit(plugin unsafe.Pointer) int { return output.FLB_ERROR } + // Optional flush timeout parameter + flushTimeoutSec, err := getConfigField(plugin, "Flush_Timeout_Sec", flushTimeoutSecDefault) + if err != nil { + log.Printf("Invalid Flush_Timeout_Sec parameter in configuration file: %s", err) + return output.FLB_ERROR + } + // Create new client client, err := getClient(ms_ctx, projectID) if err != nil { @@ -720,6 +734,7 @@ func FLBPluginInit(plugin unsafe.Pointer) int { managedStreamSlice: &streamSlice, timestampFields: timestampFields, fieldCache: buildFieldLookupCache(md), + flushTimeout: time.Duration(flushTimeoutSec) * time.Second, } // Create stream using NewManagedStream @@ -759,12 +774,15 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int } // Calls checkResponses for all streams in slice - checkAllStreamResponses(ms_ctx, &config.managedStreamSlice, false, &config.mutex, config.exactlyOnce, id) + flushCtx, flushCancel := context.WithTimeout(context.Background(), config.flushTimeout) + defer flushCancel() + + checkAllStreamResponses(flushCtx, &config.managedStreamSlice, false, &config.mutex, config.exactlyOnce, id) // Checks for need to dynamically scale - createNewStreamDynamicScaling(&config) + createNewStreamDynamicScaling(flushCtx, &config) // Create Fluent Bit decoder - dec := output.NewDecoder(data, int(length)) + dec := newDecoder(data, int(length)) // Pre-allocate binaryData slice to reduce append-driven growth (#5) binaryData := make([][]byte, 0, 256) var currsize int @@ -795,7 +813,7 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int config.mutex.Unlock() // Appending Rows - err := sendRequest(ms_ctx, binaryData, &config, leastLoadedStreamIndex) + err := sendRequest(flushCtx, binaryData, &config, leastLoadedStreamIndex) if err != nil { log.Printf("Appending data for output instance with id: %d failed in FLBPluginFlushCtx: %s", id, err) } else if config.exactlyOnce { @@ -823,7 +841,7 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int config.mutex.Unlock() // Appending Rows - err := sendRequest(ms_ctx, binaryData, &config, leastLoadedStreamIndex) + err := sendRequest(flushCtx, binaryData, &config, leastLoadedStreamIndex) if err != nil { log.Printf("Appending data for output instance with id: %d failed in FLBPluginFlushCtx: %s", id, err) } else if config.exactlyOnce { diff --git a/out_writeapi_test.go b/out_writeapi_test.go index 8102629..ba99588 100644 --- a/out_writeapi_test.go +++ b/out_writeapi_test.go @@ -446,11 +446,12 @@ func TestFLBPluginFlushCtx(t *testing.T) { } defer func() { pluginGetResult = origResultFunc }() - patchDecoder := gomonkey.ApplyFunc(output.NewDecoder, func(data unsafe.Pointer, length int) *output.FLBDecoder { + origDecoderFn := newDecoder + newDecoder = func(data unsafe.Pointer, length int) *output.FLBDecoder { checks.createDecoder++ return nil - }) - defer patchDecoder.Reset() + } + defer func() { newDecoder = origDecoderFn }() var rowSent int = 0 var rowCount int = 5 @@ -631,10 +632,11 @@ func TestFLBPluginFlushCtxDynamicScaling(t *testing.T) { } defer func() { pluginGetResult = origResultFunc }() - patchDecoder := gomonkey.ApplyFunc(output.NewDecoder, func(data unsafe.Pointer, length int) *output.FLBDecoder { + origDecoder634 := newDecoder + newDecoder = func(data unsafe.Pointer, length int) *output.FLBDecoder { return nil - }) - defer patchDecoder.Reset() + } + defer func() { newDecoder = origDecoder634 }() // Sending rows and data var rowSent int = 0 @@ -814,11 +816,12 @@ func TestFLBPluginFlushCtxExactlyOnce(t *testing.T) { } defer func() { pluginGetResult = origResultFunc }() - patchDecoder := gomonkey.ApplyFunc(output.NewDecoder, func(data unsafe.Pointer, length int) *output.FLBDecoder { + origDecoderFn := newDecoder + newDecoder = func(data unsafe.Pointer, length int) *output.FLBDecoder { checks.createDecoder++ return nil - }) - defer patchDecoder.Reset() + } + defer func() { newDecoder = origDecoderFn }() var rowSent int = 0 var rowCount int = 5 @@ -995,11 +998,12 @@ func TestFLBPluginFlushCtxErrorHandling(t *testing.T) { } defer func() { pluginGetResult = origResultFunc }() - patchDecoder := gomonkey.ApplyFunc(output.NewDecoder, func(data unsafe.Pointer, length int) *output.FLBDecoder { + origDecoderFn := newDecoder + newDecoder = func(data unsafe.Pointer, length int) *output.FLBDecoder { checks.createDecoder++ return nil - }) - defer patchDecoder.Reset() + } + defer func() { newDecoder = origDecoderFn }() var rowSent int = 0 var rowCount int = 5 From 68f07acae602e5fc792718e16b7aa652253158dd Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 18:49:40 +0900 Subject: [PATCH 05/12] perf: pool nested sub-messages (RECORD types) in dynamicpb message tree --- map_to_proto.go | 32 ++++++++--- map_to_proto_test.go | 125 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 6 deletions(-) diff --git a/map_to_proto.go b/map_to_proto.go index 9cb7826..e483960 100644 --- a/map_to_proto.go +++ b/map_to_proto.go @@ -91,8 +91,24 @@ func getPooledMessage(md protoreflect.MessageDescriptor) *dynamicpb.Message { // putPooledMessage clears all populated fields and returns the message to the pool. // Clearing via Range+Clear preserves internal map bucket memory so that // subsequent reuse avoids re-growing the map. +// Sub-messages (RECORD fields) are recursively returned to the pool before +// their parent field is cleared, enabling full reuse across the message tree. func putPooledMessage(msg *dynamicpb.Message) { - msg.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool { + msg.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { + if fd.Kind() == protoreflect.MessageKind || fd.Kind() == protoreflect.GroupKind { + if fd.IsList() { + list := v.List() + for i := 0; i < list.Len(); i++ { + if subMsg, ok := list.Get(i).Message().(*dynamicpb.Message); ok { + putPooledMessage(subMsg) + } + } + } else { + if subMsg, ok := v.Message().(*dynamicpb.Message); ok { + putPooledMessage(subMsg) + } + } + } msg.Clear(fd) return true }) @@ -156,11 +172,13 @@ func rawMapToBinary(md protoreflect.MessageDescriptor, rawData map[interface{}]i return marshalAndRelease(msg) } -// mapToMessage creates a new dynamicpb.Message and populates it from a map[string]interface{}. -// Used for nested sub-messages (which are not pooled). +// mapToMessage retrieves a pooled dynamicpb.Message and populates it from a map[string]interface{}. +// Used for nested sub-messages. The caller must not release the returned message directly; +// it will be recursively released when the top-level message is passed to marshalAndRelease. func mapToMessage(md protoreflect.MessageDescriptor, data map[string]interface{}, cache fieldLookupCache) (*dynamicpb.Message, error) { - msg := dynamicpb.NewMessage(md) + msg := getPooledMessage(md) if err := populateMessage(msg, md, data, cache); err != nil { + putPooledMessage(msg) return nil, err } return msg, nil @@ -253,8 +271,9 @@ func rawPopulateMessage(msg *dynamicpb.Message, md protoreflect.MessageDescripto } else if fd.Kind() == protoreflect.MessageKind || fd.Kind() == protoreflect.GroupKind { switch sub := val.(type) { case map[interface{}]interface{}: - subMsg := dynamicpb.NewMessage(fd.Message()) + subMsg := getPooledMessage(fd.Message()) if err := rawPopulateMessage(subMsg, fd.Message(), sub, cache); err != nil { + putPooledMessage(subMsg) return fmt.Errorf("field %q: %w", key, err) } msg.Set(fd, protoreflect.ValueOfMessage(subMsg)) @@ -354,8 +373,9 @@ func rawSetRepeatedField(msg *dynamicpb.Message, fd protoreflect.FieldDescriptor if fd.Kind() == protoreflect.MessageKind || fd.Kind() == protoreflect.GroupKind { switch sub := item.(type) { case map[interface{}]interface{}: - subMsg := dynamicpb.NewMessage(fd.Message()) + subMsg := getPooledMessage(fd.Message()) if err := rawPopulateMessage(subMsg, fd.Message(), sub, cache); err != nil { + putPooledMessage(subMsg) return fmt.Errorf("element %d: %w", i, err) } list.Append(protoreflect.ValueOfMessage(subMsg)) diff --git a/map_to_proto_test.go b/map_to_proto_test.go index 4d5daca..4fbfb37 100644 --- a/map_to_proto_test.go +++ b/map_to_proto_test.go @@ -15,6 +15,7 @@ package main import ( "encoding/base64" + "fmt" "testing" "cloud.google.com/go/bigquery/storage/apiv1/storagepb" @@ -1016,6 +1017,130 @@ func TestMessagePoolReuse_RawPath(t *testing.T) { assert.Equal(t, int64(0), msg2.Get(md.Fields().ByName("y")).Int()) // must be default } +// TestMessagePoolReuse_NestedRecord verifies that pooled sub-messages (RECORD fields) +// produce correct output across multiple sequential calls. +func TestMessagePoolReuse_NestedRecord(t *testing.T) { + schema := &storagepb.TableSchema{ + Fields: []*storagepb.TableFieldSchema{ + {Name: "id", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE}, + { + Name: "meta", + Type: storagepb.TableFieldSchema_STRUCT, + Mode: storagepb.TableFieldSchema_NULLABLE, + Fields: []*storagepb.TableFieldSchema{ + {Name: "source", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE}, + {Name: "value", Type: storagepb.TableFieldSchema_INT64, Mode: storagepb.TableFieldSchema_NULLABLE}, + }, + }, + }, + } + md := buildMD(t, schema) + cache := buildFieldLookupCache(md) + + // Call mapToBinary 100 times; sub-messages should be pooled and reused + for i := 0; i < 100; i++ { + data := map[string]interface{}{ + "id": fmt.Sprintf("row-%d", i), + "meta": map[string]interface{}{ + "source": "test", + "value": int64(i), + }, + } + b, err := mapToBinary(md, data, cache) + require.NoError(t, err) + + msg := dynamicpb.NewMessage(md) + require.NoError(t, proto.Unmarshal(b, msg)) + assert.Equal(t, fmt.Sprintf("row-%d", i), msg.Get(md.Fields().ByName("id")).String()) + metaMsg := msg.Get(md.Fields().ByName("meta")).Message() + metaMd := md.Fields().ByName("meta").Message() + assert.Equal(t, "test", metaMsg.Get(metaMd.Fields().ByName("source")).String()) + assert.Equal(t, int64(i), metaMsg.Get(metaMd.Fields().ByName("value")).Int()) + } +} + +// TestMessagePoolReuse_NoStaleFields_NestedRecord verifies that fields set inside a +// nested RECORD in one call do not leak into the next call after pool reuse. +func TestMessagePoolReuse_NoStaleFields_NestedRecord(t *testing.T) { + schema := &storagepb.TableSchema{ + Fields: []*storagepb.TableFieldSchema{ + { + Name: "info", + Type: storagepb.TableFieldSchema_STRUCT, + Mode: storagepb.TableFieldSchema_NULLABLE, + Fields: []*storagepb.TableFieldSchema{ + {Name: "a", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE}, + {Name: "b", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE}, + }, + }, + }, + } + md := buildMD(t, schema) + cache := buildFieldLookupCache(md) + infaMd := md.Fields().ByName("info").Message() + + // First call sets both nested fields + b1, err := mapToBinary(md, map[string]interface{}{ + "info": map[string]interface{}{"a": "hello", "b": "world"}, + }, cache) + require.NoError(t, err) + msg1 := dynamicpb.NewMessage(md) + require.NoError(t, proto.Unmarshal(b1, msg1)) + assert.Equal(t, "hello", msg1.Get(md.Fields().ByName("info")).Message().Get(infaMd.Fields().ByName("a")).String()) + assert.Equal(t, "world", msg1.Get(md.Fields().ByName("info")).Message().Get(infaMd.Fields().ByName("b")).String()) + + // Second call sets only "a" — "b" must NOT carry over from pooled sub-message + b2, err := mapToBinary(md, map[string]interface{}{ + "info": map[string]interface{}{"a": "only-a"}, + }, cache) + require.NoError(t, err) + msg2 := dynamicpb.NewMessage(md) + require.NoError(t, proto.Unmarshal(b2, msg2)) + assert.Equal(t, "only-a", msg2.Get(md.Fields().ByName("info")).Message().Get(infaMd.Fields().ByName("a")).String()) + assert.Equal(t, "", msg2.Get(md.Fields().ByName("info")).Message().Get(infaMd.Fields().ByName("b")).String()) +} + +// TestMessagePoolReuse_RawPath_NestedRecord verifies pool reuse for nested RECORD +// fields via the rawMapToBinary path (map[interface{}]interface{} input). +func TestMessagePoolReuse_RawPath_NestedRecord(t *testing.T) { + schema := &storagepb.TableSchema{ + Fields: []*storagepb.TableFieldSchema{ + { + Name: "nested", + Type: storagepb.TableFieldSchema_STRUCT, + Mode: storagepb.TableFieldSchema_NULLABLE, + Fields: []*storagepb.TableFieldSchema{ + {Name: "x", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE}, + {Name: "y", Type: storagepb.TableFieldSchema_INT64, Mode: storagepb.TableFieldSchema_NULLABLE}, + }, + }, + }, + } + md := buildMD(t, schema) + cache := buildFieldLookupCache(md) + nestedMd := md.Fields().ByName("nested").Message() + + // First call: set both nested fields + b1, err := rawMapToBinary(md, map[interface{}]interface{}{ + "nested": map[interface{}]interface{}{"x": []byte("first"), "y": int64(1)}, + }, cache) + require.NoError(t, err) + msg1 := dynamicpb.NewMessage(md) + require.NoError(t, proto.Unmarshal(b1, msg1)) + assert.Equal(t, "first", msg1.Get(md.Fields().ByName("nested")).Message().Get(nestedMd.Fields().ByName("x")).String()) + assert.Equal(t, int64(1), msg1.Get(md.Fields().ByName("nested")).Message().Get(nestedMd.Fields().ByName("y")).Int()) + + // Second call: set only "x" — "y" must NOT carry over from pooled sub-message + b2, err := rawMapToBinary(md, map[interface{}]interface{}{ + "nested": map[interface{}]interface{}{"x": []byte("second")}, + }, cache) + require.NoError(t, err) + msg2 := dynamicpb.NewMessage(md) + require.NoError(t, proto.Unmarshal(b2, msg2)) + assert.Equal(t, "second", msg2.Get(md.Fields().ByName("nested")).Message().Get(nestedMd.Fields().ByName("x")).String()) + assert.Equal(t, int64(0), msg2.Get(md.Fields().ByName("nested")).Message().Get(nestedMd.Fields().ByName("y")).Int()) +} + // BenchmarkRawMapToBinary benchmarks the raw map[interface{}]interface{} path. func BenchmarkRawMapToBinary(b *testing.B) { schema := mangleInputSchema(&storagepb.TableSchema{ From 05b26f61e32de72a7f7cd522b815932a99e61128 Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 19:06:07 +0900 Subject: [PATCH 06/12] perf: reuse binaryData slice across flushes to eliminate per-flush allocation binaryData ([][]byte) was re-allocated with make([][]byte, 0, 256) on every FLBPluginFlushCtx call. Move it to outputConfig so the underlying array is reused across flushes. Nil out retained elements before each reset so the GC can reclaim the previous batch's byte slices. --- out_writeapi.go | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/out_writeapi.go b/out_writeapi.go index c037ed1..ee9d098 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -64,6 +64,7 @@ type outputConfig struct { timestampFields []string fieldCache fieldLookupCache flushTimeout time.Duration + binaryData [][]byte // reused across flushes to avoid per-flush allocation } var ( @@ -783,8 +784,12 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int // Create Fluent Bit decoder dec := newDecoder(data, int(length)) - // Pre-allocate binaryData slice to reduce append-driven growth (#5) - binaryData := make([][]byte, 0, 256) + // Reuse binaryData slice across flushes to avoid per-flush allocation. + // Nil out retained elements before reset so GC can reclaim previous batch's bytes. + for i := range config.binaryData { + config.binaryData[i] = nil + } + binaryData := config.binaryData[:0] var currsize int // Keeps track of the number of rows previously sent var rowCounter int64 @@ -824,7 +829,10 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int rowCounter = 0 - // Reuse the underlying array instead of nil to avoid re-allocation + // Nil out sent elements so GC can reclaim them, then reset slice. + for i := range binaryData { + binaryData[i] = nil + } binaryData = binaryData[:0] currsize = 0 @@ -850,6 +858,9 @@ func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int config.mutex.Unlock() } + // Write back the grown slice so its capacity is reused on the next flush. + config.binaryData = binaryData + return output.FLB_OK } From 4482145c13f7ccc24d8b3c6cb5ffcced9ee0a890 Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 19:13:24 +0900 Subject: [PATCH 07/12] fix: drain pending AppendRows responses before closing streams on exit FLBPluginExitCtx called checkAllStreamResponses with waitForResponse=false, which skips any in-flight response that has not yet arrived. This could cause data loss if AppendRows results were still pending when the plugin exited. Fix: pass waitForResponse=true so pluginGetResult blocks on each queued result, fully draining the queue before finalizeCloseAllStreams runs. Also switch FLBPluginExitCtx to use the injectable getFLBPluginContext variable (consistent with FLBPluginFlushCtx) to enable unit testing without gomonkey. Add TestFLBPluginExitCtxDrainsPendingResponses to verify that responses marked as not-yet-ready are still drained during exit. Fix MockManagedStream.Close and MockManagedWriterClient.Close to use their CloseFunc fields when set, avoiding nil-pointer panics in tests that call ExitCtx. --- out_writeapi.go | 8 +-- out_writeapi_test.go | 120 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 4 deletions(-) diff --git a/out_writeapi.go b/out_writeapi.go index ee9d098..824e552 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -873,7 +873,7 @@ func FLBPluginExit() int { //export FLBPluginExitCtx func FLBPluginExitCtx(ctx unsafe.Pointer) int { // Get context - id := output.FLBPluginGetContext(ctx).(int) + id := getFLBPluginContext(ctx) // Locate stream in map config, ok := configMap[id] @@ -882,9 +882,9 @@ func FLBPluginExitCtx(ctx unsafe.Pointer) int { return output.FLB_ERROR } - // Calls checkResponses, finalizes, and closes each stream - // If there is an error it preserves it and then continues to close and finalize all other streams - checkAllStreamResponses(ms_ctx, &config.managedStreamSlice, false, &config.mutex, config.exactlyOnce, id) + // Drain all pending responses before closing streams to avoid data loss. + // waitForResponse=true blocks until every in-flight AppendRows result is received. + checkAllStreamResponses(ms_ctx, &config.managedStreamSlice, true, &config.mutex, config.exactlyOnce, id) errFlag := finalizeCloseAllStreams(&config, id) if config.client != nil { diff --git a/out_writeapi_test.go b/out_writeapi_test.go index ba99588..ace2af7 100644 --- a/out_writeapi_test.go +++ b/out_writeapi_test.go @@ -98,6 +98,9 @@ func (m *MockManagedWriterClient) GetWriteStream(ctx context.Context, req *stora } func (m *MockManagedWriterClient) Close() error { + if m.CloseFunc != nil { + return m.CloseFunc() + } return m.client.Close() } @@ -327,6 +330,9 @@ func (m *MockManagedStream) StreamName() string { } func (m *MockManagedStream) Close() error { + if m.CloseFunc != nil { + return m.CloseFunc() + } return m.managedstream.Close() } @@ -1050,6 +1056,120 @@ func TestFLBPluginFlushCtxErrorHandling(t *testing.T) { assert.Equal(t, expectGotRecord, checks.gotRecord) } +// TestFLBPluginExitCtxDrainsPendingResponses verifies that FLBPluginExitCtx +// drains all pending (not-yet-ready) AppendRows responses before closing streams. +func TestFLBPluginExitCtxDrainsPendingResponses(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 }() + + var appendRowsCalled int + mockMS := &MockManagedStream{ + AppendRowsFunc: func(ctx context.Context, data [][]byte, opts ...managedwriter.AppendOption) (*managedwriter.AppendResult, error) { + appendRowsCalled++ + return nil, nil // nil AppendResult queued; drained via mocked pluginGetResult + }, + 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 }() + + // isReady always returns false — responses are "in-flight" and not yet ready. + origIsReady := isReady + isReady = func(_ *managedwriter.AppendResult) bool { return false } + defer func() { isReady = origIsReady }() + + // pluginGetResult counts how many pending responses were drained. + var drainCount int + origGetResult := pluginGetResult + pluginGetResult = func(_ *managedwriter.AppendResult, _ context.Context) (int64, error) { + drainCount++ + return -1, nil + } + defer func() { pluginGetResult = origGetResult }() + + // Send one row via flush so AppendRows is called and a result is queued. + 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) + assert.Equal(t, 1, appendRowsCalled) // one AppendRows call queued a result + + // After flush, the queue holds 1 pending response (isReady=false so flush didn't drain it). + // FLBPluginExitCtx must drain it with waitForResponse=true. + assert.Equal(t, 0, drainCount) // not yet drained + + exitResult := FLBPluginExitCtx(pointerValue) + assert.Equal(t, output.FLB_OK, exitResult) + assert.Equal(t, 1, drainCount) // drained during exit +} + // TestConvertTimestampFieldsRaw tests the convertTimestampFieldsRaw function // for raw Fluent Bit records (map[interface{}]interface{}). func TestConvertTimestampFieldsRaw(t *testing.T) { From 5bcb910bb13d388e43e87a439f94abb7cbb310a2 Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 19:15:04 +0900 Subject: [PATCH 08/12] Flush_Timeout_Sec default 30 --- out_writeapi.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/out_writeapi.go b/out_writeapi.go index 824e552..c325ec6 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -84,7 +84,7 @@ const ( minQueueRequests = 10 dateTimeDefault = true maxUnixSeconds = 4102444800 // 2100-01-01 00:00:00 UTC in seconds - flushTimeoutSecDefault = 60 + flushTimeoutSecDefault = 30 ) // This function mangles the top-level and complex (struct) BigQuery schema to convert NUMERIC, BIGNUMERIC, DATETIME, TIME, and JSON fields to STRING. From 471731511e30dad78e7dc0efa24d5f2b07f8f122 Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 19:22:02 +0900 Subject: [PATCH 09/12] docs: add comments to magic number constants and fix typo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Annotate const block (out_writeapi.go) with explanations for chunkSizeLimit, queueRequestDefault, queueByteDefault, queueRequestScalingPercent, maxNumStreamsPerInstance, minQueueRequests, maxUnixSeconds, and flushTimeoutSecDefault. - Fix "Unsuccesful" → "Unsuccessful" in sendRequestRetries comment. --- out_writeapi.go | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/out_writeapi.go b/out_writeapi.go index c325ec6..9a916be 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -74,17 +74,17 @@ var ( ) const ( - chunkSizeLimit = 9 * 1024 * 1024 - queueRequestDefault = 1000 - queueByteDefault = 100 * 1024 * 1024 + chunkSizeLimit = 9 * 1024 * 1024 // BigQuery Storage Write API AppendRows hard limit (10MB minus overhead) + queueRequestDefault = 1000 // default Max_Queue_Requests (max in-flight AppendRows per stream) + queueByteDefault = 100 * 1024 * 1024 // default Max_Queue_Bytes: 100 MB exactlyOnceDefault = false - queueRequestScalingPercent = 0.8 + queueRequestScalingPercent = 0.8 // queue utilization threshold (80%) above which a new stream is created numRetriesDefault = 4 - maxNumStreamsPerInstance = 10 - minQueueRequests = 10 + maxNumStreamsPerInstance = 10 // upper bound on dynamic stream count per output instance + minQueueRequests = 10 // minimum Max_Queue_Requests accepted (prevents starvation during scaling) dateTimeDefault = true - maxUnixSeconds = 4102444800 // 2100-01-01 00:00:00 UTC in seconds - flushTimeoutSecDefault = 30 + maxUnixSeconds = 4102444800 // 2100-01-01 00:00:00 UTC; values above this are treated as already in microseconds + flushTimeoutSecDefault = 30 // default Flush_Timeout_Sec: 30 seconds per flush call ) // This function mangles the top-level and complex (struct) BigQuery schema to convert NUMERIC, BIGNUMERIC, DATETIME, TIME, and JSON fields to STRING. @@ -404,7 +404,7 @@ func sendRequestRetries(ctx context.Context, data [][]byte, config **outputConfi if err == nil { break } - // Unsuccesful data append + // Unsuccessful data append if rebuildPredicate(err) { currStream.managedstream.Finalize(ctx) currStream.managedstream.Close() From e14ef73b4c7eb4f64927f833b37fb28d0c387d6e Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 19:31:08 +0900 Subject: [PATCH 10/12] perf: atomic configID, GC-safe response queue drain, and zero-alloc []byte parsing --- map_to_proto.go | 56 +++++++++++++++++++++++++++++++++++++++++-------- out_writeapi.go | 42 ++++++++++++++++++++----------------- 2 files changed, 70 insertions(+), 28 deletions(-) diff --git a/map_to_proto.go b/map_to_proto.go index e483960..fd7f858 100644 --- a/map_to_proto.go +++ b/map_to_proto.go @@ -18,6 +18,7 @@ import ( "fmt" "strconv" "sync" + "unsafe" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/reflect/protoreflect" @@ -434,7 +435,7 @@ func toProtoString(val interface{}) (protoreflect.Value, error) { case string: return protoreflect.ValueOfString(v), nil case []byte: - return protoreflect.ValueOfString(string(v)), nil + return protoreflect.ValueOfString(unsafe.String(unsafe.SliceData(v), len(v))), nil case int: return protoreflect.ValueOfString(strconv.Itoa(v)), nil case int64: @@ -478,7 +479,16 @@ func toProtoInt64(val interface{}) (protoreflect.Value, error) { } return protoreflect.ValueOfInt64(i), nil case []byte: - return toProtoInt64(string(v)) + s := unsafe.String(unsafe.SliceData(v), len(v)) + i, err := strconv.ParseInt(s, 10, 64) + if err != nil { + f, ferr := strconv.ParseFloat(s, 64) + if ferr != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to int64: %w", v, err) + } + return protoreflect.ValueOfInt64(int64(f)), nil + } + return protoreflect.ValueOfInt64(i), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to int64", val) } @@ -507,7 +517,11 @@ func toProtoInt32(val interface{}) (protoreflect.Value, error) { } return protoreflect.ValueOfInt32(int32(i)), nil case []byte: - return toProtoInt32(string(v)) + i, err := strconv.ParseInt(unsafe.String(unsafe.SliceData(v), len(v)), 10, 32) + if err != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to int32: %w", v, err) + } + return protoreflect.ValueOfInt32(int32(i)), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to int32", val) } @@ -532,7 +546,11 @@ func toProtoUint64(val interface{}) (protoreflect.Value, error) { } return protoreflect.ValueOfUint64(u), nil case []byte: - return toProtoUint64(string(v)) + u, err := strconv.ParseUint(unsafe.String(unsafe.SliceData(v), len(v)), 10, 64) + if err != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to uint64: %w", v, err) + } + return protoreflect.ValueOfUint64(u), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to uint64", val) } @@ -557,7 +575,11 @@ func toProtoUint32(val interface{}) (protoreflect.Value, error) { } return protoreflect.ValueOfUint32(uint32(u)), nil case []byte: - return toProtoUint32(string(v)) + u, err := strconv.ParseUint(unsafe.String(unsafe.SliceData(v), len(v)), 10, 32) + if err != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to uint32: %w", v, err) + } + return protoreflect.ValueOfUint32(uint32(u)), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to uint32", val) } @@ -582,7 +604,11 @@ func toProtoDouble(val interface{}) (protoreflect.Value, error) { } return protoreflect.ValueOfFloat64(f), nil case []byte: - return toProtoDouble(string(v)) + f, err := strconv.ParseFloat(unsafe.String(unsafe.SliceData(v), len(v)), 64) + if err != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to float64: %w", v, err) + } + return protoreflect.ValueOfFloat64(f), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to float64", val) } @@ -605,7 +631,11 @@ func toProtoFloat(val interface{}) (protoreflect.Value, error) { } return protoreflect.ValueOfFloat32(float32(f)), nil case []byte: - return toProtoFloat(string(v)) + f, err := strconv.ParseFloat(unsafe.String(unsafe.SliceData(v), len(v)), 32) + if err != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to float32: %w", v, err) + } + return protoreflect.ValueOfFloat32(float32(f)), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to float32", val) } @@ -630,7 +660,11 @@ func toProtoBool(val interface{}) (protoreflect.Value, error) { case float64: return protoreflect.ValueOfBool(v != 0), nil case []byte: - return toProtoBool(string(v)) + b, err := strconv.ParseBool(unsafe.String(unsafe.SliceData(v), len(v))) + if err != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to bool: %w", v, err) + } + return protoreflect.ValueOfBool(b), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to bool", val) } @@ -676,7 +710,11 @@ func toProtoEnum(val interface{}) (protoreflect.Value, error) { } return protoreflect.ValueOfEnum(protoreflect.EnumNumber(i)), nil case []byte: - return toProtoEnum(string(v)) + i, err := strconv.ParseInt(unsafe.String(unsafe.SliceData(v), len(v)), 10, 32) + if err != nil { + return protoreflect.Value{}, fmt.Errorf("cannot convert []byte %q to enum: %w", v, err) + } + return protoreflect.ValueOfEnum(protoreflect.EnumNumber(i)), nil default: return protoreflect.Value{}, fmt.Errorf("cannot convert %T to enum", val) } diff --git a/out_writeapi.go b/out_writeapi.go index 9a916be..23b052f 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -23,6 +23,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "unsafe" @@ -70,7 +71,7 @@ type outputConfig struct { var ( ms_ctx = context.Background() configMap = make(map[int]*outputConfig) - configID = 0 + configID atomic.Int64 ) const ( @@ -286,25 +287,30 @@ var pluginGetResult = func(result *managedwriter.AppendResult, ctx context.Conte // And wait for the next ready response from WriteAPI // This function returns an int which is the length of the queue after being checked or -1 if an error occured func checkResponses(curr_ctx context.Context, streamSlice *[]*streamConfig, waitForResponse bool, exactlyOnceConf bool, id int, streamIndex int) int { - currQueuePointer := (*streamSlice)[streamIndex].appendResults - for len(*currQueuePointer) > 0 { - if exactlyOnceConf { - log.Printf("Asynchronous response queue has non-zero size when exactly-once is configured") - break - } - queueHead := (*currQueuePointer)[0] - if waitForResponse || isReady(queueHead) { - _, err := pluginGetResult(queueHead, curr_ctx) - *currQueuePointer = (*currQueuePointer)[1:] + currQueue := (*streamSlice)[streamIndex].appendResults + if exactlyOnceConf && len(*currQueue) > 0 { + log.Printf("Asynchronous response queue has non-zero size when exactly-once is configured") + return len(*currQueue) + } + writeIdx := 0 + for readIdx := 0; readIdx < len(*currQueue); readIdx++ { + result := (*currQueue)[readIdx] + if waitForResponse || isReady(result) { + _, err := pluginGetResult(result, curr_ctx) if err != nil { log.Printf("Encountered error:%s while verifying the server response to a data append for output instance with id: %d", err, id) } } else { - break + (*currQueue)[writeIdx] = result + writeIdx++ } - } - return len(*currQueuePointer) + // nil out drained slots to allow GC to reclaim AppendResult pointers + for i := writeIdx; i < len(*currQueue); i++ { + (*currQueue)[i] = nil + } + *currQueue = (*currQueue)[:writeIdx] + return writeIdx } // This function checks the responses for all streams in the slice for each instance @@ -746,13 +752,11 @@ func FLBPluginInit(plugin unsafe.Pointer) int { return output.FLB_ERROR } - configMap[configID] = &config + id := int(configID.Add(1) - 1) + configMap[id] = &config // Creating FLB context for each output, enables multiinstancing - config.mutex.Lock() - output.FLBPluginSetContext(plugin, configID) - configID = configID + 1 - config.mutex.Unlock() + output.FLBPluginSetContext(plugin, id) return output.FLB_OK } From 1c55fa672d479a1d63524bb1db1397761ce89b1f Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 19:37:15 +0900 Subject: [PATCH 11/12] =?UTF-8?q?refactor:=20extract=20toSubMessage=20help?= =?UTF-8?q?er=20to=20eliminate=20duplicate=20map=E2=86=92message=20convers?= =?UTF-8?q?ion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add toSubMessage(val, md, cache) as the single authoritative path for converting either map[string]interface{} or map[interface{}]interface{} values into a populated *dynamicpb.Message. Replace the duplicated type-switch/assertion blocks in populateMessage, rawPopulateMessage, setRepeatedField, and rawSetRepeatedField with calls to the new helper. rawSetRepeatedField becomes identical to setRepeatedField after this change and is therefore deleted; its one call site in rawPopulateMessage now calls setRepeatedField directly. --- map_to_proto.go | 99 ++++++++++++++----------------------------------- 1 file changed, 27 insertions(+), 72 deletions(-) diff --git a/map_to_proto.go b/map_to_proto.go index fd7f858..72f1f59 100644 --- a/map_to_proto.go +++ b/map_to_proto.go @@ -185,6 +185,26 @@ func mapToMessage(md protoreflect.MessageDescriptor, data map[string]interface{} return msg, nil } +// toSubMessage converts val (map[string]interface{} or map[interface{}]interface{}) +// into a populated *dynamicpb.Message using a pooled message. +// It is the single authoritative path for building sub-messages from either the +// typed (mapToBinary) or raw (rawMapToBinary) Fluent Bit record paths. +func toSubMessage(val interface{}, md protoreflect.MessageDescriptor, cache fieldLookupCache) (*dynamicpb.Message, error) { + switch sub := val.(type) { + case map[interface{}]interface{}: + subMsg := getPooledMessage(md) + if err := rawPopulateMessage(subMsg, md, sub, cache); err != nil { + putPooledMessage(subMsg) + return nil, err + } + return subMsg, nil + case map[string]interface{}: + return mapToMessage(md, sub, cache) + default: + return nil, fmt.Errorf("expected map for message field, got %T", val) + } +} + // populateMessage fills an existing dynamicpb.Message from a map[string]interface{}. // Extracted from mapToMessage so that callers can supply a pooled message func populateMessage(msg *dynamicpb.Message, md protoreflect.MessageDescriptor, data map[string]interface{}, cache fieldLookupCache) error { @@ -215,11 +235,7 @@ func populateMessage(msg *dynamicpb.Message, md protoreflect.MessageDescriptor, return fmt.Errorf("field %q: %w", key, err) } } else if fd.Kind() == protoreflect.MessageKind || fd.Kind() == protoreflect.GroupKind { - subMap, ok := val.(map[string]interface{}) - if !ok { - return fmt.Errorf("field %q: expected map for message field, got %T", key, val) - } - subMsg, err := mapToMessage(fd.Message(), subMap, cache) + subMsg, err := toSubMessage(val, fd.Message(), cache) if err != nil { return fmt.Errorf("field %q: %w", key, err) } @@ -266,27 +282,15 @@ func rawPopulateMessage(msg *dynamicpb.Message, md protoreflect.MessageDescripto } if fd.IsList() { - if err := rawSetRepeatedField(msg, fd, val, cache); err != nil { + if err := setRepeatedField(msg, fd, val, cache); err != nil { return fmt.Errorf("field %q: %w", key, err) } } else if fd.Kind() == protoreflect.MessageKind || fd.Kind() == protoreflect.GroupKind { - switch sub := val.(type) { - case map[interface{}]interface{}: - subMsg := getPooledMessage(fd.Message()) - if err := rawPopulateMessage(subMsg, fd.Message(), sub, cache); err != nil { - putPooledMessage(subMsg) - return fmt.Errorf("field %q: %w", key, err) - } - msg.Set(fd, protoreflect.ValueOfMessage(subMsg)) - case map[string]interface{}: - subMsg, err := mapToMessage(fd.Message(), sub, cache) - if err != nil { - return fmt.Errorf("field %q: %w", key, err) - } - msg.Set(fd, protoreflect.ValueOfMessage(subMsg)) - default: - return fmt.Errorf("field %q: expected map for message field, got %T", key, val) + subMsg, err := toSubMessage(val, fd.Message(), cache) + if err != nil { + return fmt.Errorf("field %q: %w", key, err) } + msg.Set(fd, protoreflect.ValueOfMessage(subMsg)) } else { pv, err := goToProtoScalar(fd, val) if err != nil { @@ -335,11 +339,7 @@ func setRepeatedField(msg *dynamicpb.Message, fd protoreflect.FieldDescriptor, v continue } if fd.Kind() == protoreflect.MessageKind || fd.Kind() == protoreflect.GroupKind { - subMap, ok := item.(map[string]interface{}) - if !ok { - return fmt.Errorf("element %d: expected map for repeated message, got %T", i, item) - } - subMsg, err := mapToMessage(fd.Message(), subMap, cache) + subMsg, err := toSubMessage(item, fd.Message(), cache) if err != nil { return fmt.Errorf("element %d: %w", i, err) } @@ -355,51 +355,6 @@ func setRepeatedField(msg *dynamicpb.Message, fd protoreflect.FieldDescriptor, v return nil } -// rawSetRepeatedField populates a repeated proto field from a raw slice, -// handling map[interface{}]interface{} elements for nested messages -func rawSetRepeatedField(msg *dynamicpb.Message, fd protoreflect.FieldDescriptor, val interface{}, cache fieldLookupCache) error { - slice, ok := val.([]interface{}) - if !ok { - return fmt.Errorf("expected []interface{} for repeated field, got %T", val) - } - if len(slice) == 0 { - return nil - } - - list := msg.Mutable(fd).List() - for i, item := range slice { - if item == nil { - continue - } - if fd.Kind() == protoreflect.MessageKind || fd.Kind() == protoreflect.GroupKind { - switch sub := item.(type) { - case map[interface{}]interface{}: - subMsg := getPooledMessage(fd.Message()) - if err := rawPopulateMessage(subMsg, fd.Message(), sub, cache); err != nil { - putPooledMessage(subMsg) - return fmt.Errorf("element %d: %w", i, err) - } - list.Append(protoreflect.ValueOfMessage(subMsg)) - case map[string]interface{}: - subMsg, err := mapToMessage(fd.Message(), sub, cache) - if err != nil { - return fmt.Errorf("element %d: %w", i, err) - } - list.Append(protoreflect.ValueOfMessage(subMsg)) - default: - return fmt.Errorf("element %d: expected map for repeated message, got %T", i, item) - } - } else { - pv, err := goToProtoScalar(fd, item) - if err != nil { - return fmt.Errorf("element %d: %w", i, err) - } - list.Append(pv) - } - } - return nil -} - // goToProtoScalar converts a Go interface{} value to the appropriate // protoreflect.Value for the given field descriptor's kind. // Handles all proto scalar types that appear in BigQuery Storage Write API schemas. From 211f4e9058a3964002ae8868119b32d235e6386e Mon Sep 17 00:00:00 2001 From: yuzone Date: Fri, 27 Mar 2026 19:54:59 +0900 Subject: [PATCH 12/12] refactor: simplify **outputConfig to *outputConfig and split FLBPluginFlushCtx Change all internal functions (buildStream, sendRequest*, createNew- StreamDynamicScaling, finalizeCloseAllStreams) from **outputConfig to *outputConfig, eliminating triple-dereference patterns throughout. Remove the configPointer intermediate variable in FLBPluginInit. Convert buildStream to an injectable var (consistent with newDecoder, getClient, etc.) and migrate TestFLBPluginFlushCtxDynamicScaling off gomonkey.ApplyFunc to fix ARM64 reliability. Extract flushChunk (stream selection + send + offset update) and decodeAndSerializeRecords (record loop + mid-chunk flush) from FLBPluginFlushCtx, reducing it from 95 to 34 lines. --- out_writeapi.go | 223 +++++++++++++++++++++---------------------- out_writeapi_test.go | 23 ++--- 2 files changed, 120 insertions(+), 126 deletions(-) diff --git a/out_writeapi.go b/out_writeapi.go index 23b052f..2712620 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -348,21 +348,21 @@ func getConfigField[T int | bool](plugin unsafe.Pointer, key string, defaultval return finval, nil } -// This function creates a new managed stream based on the config struct fields -func buildStream(ctx context.Context, config **outputConfig, streamIndex int) error { - currManagedStream, err := getWriter((*config).client, ctx, (*config).currProjectID, - managedwriter.WithType((*config).streamType), - managedwriter.WithDestinationTable((*config).tableRef), +// buildStream is a var so tests can replace it without gomonkey. +var buildStream = func(ctx context.Context, config *outputConfig, streamIndex int) error { + currManagedStream, err := getWriter(config.client, ctx, config.currProjectID, + managedwriter.WithType(config.streamType), + managedwriter.WithDestinationTable(config.tableRef), // Use the descriptor proto when creating the new managed stream - managedwriter.WithSchemaDescriptor((*config).schemaDesc), - managedwriter.EnableWriteRetries((*config).enableRetry), - managedwriter.WithMaxInflightBytes((*config).maxQueueBytes), - managedwriter.WithMaxInflightRequests((*config).maxQueueRequests), + managedwriter.WithSchemaDescriptor(config.schemaDesc), + managedwriter.EnableWriteRetries(config.enableRetry), + managedwriter.WithMaxInflightBytes(config.maxQueueBytes), + managedwriter.WithMaxInflightRequests(config.maxQueueRequests), managedwriter.WithDefaultMissingValueInterpretation(storagepb.AppendRowsRequest_DEFAULT_VALUE), managedwriter.WithTraceID("FluentBit"), ) - streamSlice := *(*config).managedStreamSlice + streamSlice := *config.managedStreamSlice if err == nil { (streamSlice)[streamIndex].managedstream = currManagedStream @@ -382,11 +382,11 @@ func rebuildPredicate(err error) bool { } // This function sends and checks the responses for data through a committed stream with exactly once functionality -func sendRequestExactlyOnce(ctx context.Context, data [][]byte, config **outputConfig, streamIndex int) error { - (*config).mutex.Lock() - defer (*config).mutex.Unlock() +func sendRequestExactlyOnce(ctx context.Context, data [][]byte, config *outputConfig, streamIndex int) error { + config.mutex.Lock() + defer config.mutex.Unlock() - currStream := (*(*config).managedStreamSlice)[streamIndex] + currStream := (*config.managedStreamSlice)[streamIndex] appendResult, err := currStream.managedstream.AppendRows(ctx, data, managedwriter.WithOffset(currStream.offsetCounter)) if err != nil { @@ -401,10 +401,10 @@ func sendRequestExactlyOnce(ctx context.Context, data [][]byte, config **outputC } // This function enables synchronous retries and rebuilding a valid stream based on the server response -func sendRequestRetries(ctx context.Context, data [][]byte, config **outputConfig, streamIndex int) error { - retryer := newStatelessRetryer((*config).numRetries) +func sendRequestRetries(ctx context.Context, data [][]byte, config *outputConfig, streamIndex int) error { + retryer := newStatelessRetryer(config.numRetries) attempt := 0 - currStream := (*(*config).managedStreamSlice)[streamIndex] + currStream := (*config.managedStreamSlice)[streamIndex] for { err := sendRequestExactlyOnce(ctx, data, config, streamIndex) if err == nil { @@ -435,10 +435,10 @@ func sendRequestRetries(ctx context.Context, data [][]byte, config **outputConfi } // This function sends data and appends the responses to a queue to be checked asynchronously through a default stream with at least once functionality -func sendRequestDefault(ctx context.Context, data [][]byte, config **outputConfig, streamIndex int) error { - (*config).mutex.Lock() - defer (*config).mutex.Unlock() - currStream := (*(*config).managedStreamSlice)[streamIndex] +func sendRequestDefault(ctx context.Context, data [][]byte, config *outputConfig, streamIndex int) error { + config.mutex.Lock() + defer config.mutex.Unlock() + currStream := (*config.managedStreamSlice)[streamIndex] appendResult, err := currStream.managedstream.AppendRows(ctx, data) if err != nil { @@ -450,9 +450,9 @@ func sendRequestDefault(ctx context.Context, data [][]byte, config **outputConfi } // This function cases on the exactly/at-least once functionality and sends the data accordingly -func sendRequest(ctx context.Context, data [][]byte, config **outputConfig, streamIndex int) error { +func sendRequest(ctx context.Context, data [][]byte, config *outputConfig, streamIndex int) error { if len(data) > 0 { - if (*config).exactlyOnce { + if config.exactlyOnce { return sendRequestRetries(ctx, data, config, streamIndex) } else { return sendRequestDefault(ctx, data, config, streamIndex) @@ -497,13 +497,13 @@ 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) { - (*config).mutex.Lock() - defer (*config).mutex.Unlock() - if len(*(*config).managedStreamSlice) < maxNumStreamsPerInstance { +func createNewStreamDynamicScaling(ctx context.Context, config *outputConfig) { + config.mutex.Lock() + defer config.mutex.Unlock() + if len(*config.managedStreamSlice) < maxNumStreamsPerInstance { // Gets stream with least values in queue - mostEfficient := getLeastLoadedStream((*config).managedStreamSlice) - mostEfficientQueueLength := len(*(*(*config).managedStreamSlice)[mostEfficient].appendResults) + mostEfficient := getLeastLoadedStream(config.managedStreamSlice) + mostEfficientQueueLength := len(*(*config.managedStreamSlice)[mostEfficient].appendResults) var newResQueue []*managedwriter.AppendResult var newStream = streamConfig{ @@ -511,14 +511,14 @@ func createNewStreamDynamicScaling(ctx context.Context, config **outputConfig) { appendResults: &newResQueue, } - if mostEfficientQueueLength > (*config).requestCountThreshold { - *(*config).managedStreamSlice = append(*(*config).managedStreamSlice, &newStream) - newStreamIndex := len(*(*config).managedStreamSlice) - 1 + if mostEfficientQueueLength > config.requestCountThreshold { + *config.managedStreamSlice = append(*config.managedStreamSlice, &newStream) + newStreamIndex := len(*config.managedStreamSlice) - 1 err := buildStream(ctx, config, newStreamIndex) if err != nil { - log.Printf("Creating an additional managed stream with destination table: %s failed in FLBPluginInit: %s", (*config).tableRef, err) + 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 - *(*config).managedStreamSlice = (*(*config).managedStreamSlice)[:newStreamIndex] + *config.managedStreamSlice = (*config.managedStreamSlice)[:newStreamIndex] } } @@ -598,14 +598,14 @@ var getFLBPluginContext = func(ctx unsafe.Pointer) int { } // Finalizes and Closes all streams in slice for a given instance -func finalizeCloseAllStreams(config **outputConfig, id int) bool { - (*config).mutex.Lock() - defer (*config).mutex.Unlock() +func finalizeCloseAllStreams(config *outputConfig, id int) bool { + config.mutex.Lock() + defer config.mutex.Unlock() errFlag := false - streamSlice := (*config).managedStreamSlice - for i := 0; i < len(*(*config).managedStreamSlice); i++ { + streamSlice := config.managedStreamSlice + for i := 0; i < len(*config.managedStreamSlice); i++ { if (*streamSlice)[i].managedstream != nil { - if (*config).exactlyOnce { + if config.exactlyOnce { if _, err := (*streamSlice)[i].managedstream.Finalize(ms_ctx); err != nil { log.Printf("Finalizing managed stream for output instance with id %d and stream index %d failed in FLBPluginExit: %s", id, i, err) errFlag = true @@ -745,8 +745,7 @@ func FLBPluginInit(plugin unsafe.Pointer) int { } // Create stream using NewManagedStream - configPointer := &config - err = buildStream(ms_ctx, &configPointer, 0) + err = buildStream(ms_ctx, &config, 0) if err != nil { log.Printf("Creating a new managed stream with destination table: %s failed in FLBPluginInit: %s", tableReference, err) return output.FLB_ERROR @@ -767,100 +766,94 @@ func FLBPluginFlush(data unsafe.Pointer, length C.int, tag *C.char) int { return output.FLB_OK } -//export FLBPluginFlushCtx -func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int { - id := getFLBPluginContext(ctx) - // Locate stream in map - // Look up through reference - config, ok := configMap[id] - if !ok { - log.Printf("Finding configuration for output instance with id: %d failed in FLBPluginFlushCtx", id) - return output.FLB_ERROR - } - - // Calls checkResponses for all streams in slice - flushCtx, flushCancel := context.WithTimeout(context.Background(), config.flushTimeout) - defer flushCancel() - - checkAllStreamResponses(flushCtx, &config.managedStreamSlice, false, &config.mutex, config.exactlyOnce, id) - // Checks for need to dynamically scale - createNewStreamDynamicScaling(flushCtx, &config) +// flushChunk picks the least-loaded stream, sends binaryData, and (for +// exactly-once mode) increments the stream's offset counter by rowCount. +func flushChunk(ctx context.Context, config *outputConfig, id int, binaryData [][]byte, rowCount int64) { + config.mutex.Lock() + streamIndex := getLeastLoadedStream(config.managedStreamSlice) + config.mutex.Unlock() - // Create Fluent Bit decoder - dec := newDecoder(data, int(length)) - // Reuse binaryData slice across flushes to avoid per-flush allocation. - // Nil out retained elements before reset so GC can reclaim previous batch's bytes. - for i := range config.binaryData { - config.binaryData[i] = nil + if err := sendRequest(ctx, binaryData, config, streamIndex); err != nil { + log.Printf("Appending data for output instance with id: %d failed in FLBPluginFlushCtx: %s", id, err) + return } - binaryData := config.binaryData[:0] + if config.exactlyOnce { + config.mutex.Lock() + (*config.managedStreamSlice)[streamIndex].offsetCounter += rowCount + config.mutex.Unlock() + } +} + +// decodeAndSerializeRecords decodes Fluent Bit records from dec, serializes each +// to proto binary, and accumulates them in binaryData. When the accumulated size +// reaches config.maxChunkSize, it calls flushChunk and resets binaryData +// (preserving capacity for reuse). It returns the remaining (unsent) binaryData +// and the row count for the final chunk. +func decodeAndSerializeRecords(ctx context.Context, config *outputConfig, id int, dec *output.FLBDecoder, binaryData [][]byte) ([][]byte, int64) { var currsize int - // Keeps track of the number of rows previously sent var rowCounter int64 - // Iterate Records for { - // Extract Record ret, _, record := output.GetRecord(dec) if ret != 0 { break } - // Convert timestamp fields in-place on the raw record convertTimestampFieldsRaw(record, config.timestampFields) - // Serialize data directly from raw Fluent Bit record buf, err := rawMapToBinary(config.messageDescriptor, record, config.fieldCache) if err != nil { log.Printf("Transforming row with value:%v from map to binary data for output instance with id: %d failed in FLBPluginFlushCtx: %s", record, id, err) - } else { - // Successful data transformation - if (currsize + len(buf)) >= config.maxChunkSize { - // Re-evaluate the least loaded stream per chunk to distribute load across streams. - config.mutex.Lock() - leastLoadedStreamIndex := getLeastLoadedStream(config.managedStreamSlice) - config.mutex.Unlock() - - // Appending Rows - err := sendRequest(flushCtx, binaryData, &config, leastLoadedStreamIndex) - if err != nil { - log.Printf("Appending data for output instance with id: %d failed in FLBPluginFlushCtx: %s", id, err) - } else if config.exactlyOnce { - config.mutex.Lock() - (*config.managedStreamSlice)[leastLoadedStreamIndex].offsetCounter += rowCounter - config.mutex.Unlock() - } - - rowCounter = 0 - - // Nil out sent elements so GC can reclaim them, then reset slice. - for i := range binaryData { - binaryData[i] = nil - } - binaryData = binaryData[:0] - currsize = 0 + continue + } + if (currsize + len(buf)) >= config.maxChunkSize { + flushChunk(ctx, config, id, binaryData, rowCounter) + rowCounter = 0 + // Nil out sent elements so GC can reclaim them, then reset slice. + for i := range binaryData { + binaryData[i] = nil } - binaryData = append(binaryData, buf) - // Include the protobuf overhead to the currsize variable - currsize += (len(buf) + 2) - rowCounter++ + binaryData = binaryData[:0] + currsize = 0 } + binaryData = append(binaryData, buf) + // Include the protobuf overhead in the size estimate. + currsize += (len(buf) + 2) + rowCounter++ } - // Re-evaluate the least loaded stream for the final chunk. - config.mutex.Lock() - leastLoadedStreamIndex := getLeastLoadedStream(config.managedStreamSlice) - config.mutex.Unlock() - // Appending Rows - err := sendRequest(flushCtx, binaryData, &config, leastLoadedStreamIndex) - if err != nil { - log.Printf("Appending data for output instance with id: %d failed in FLBPluginFlushCtx: %s", id, err) - } else if config.exactlyOnce { - config.mutex.Lock() - (*config.managedStreamSlice)[leastLoadedStreamIndex].offsetCounter += rowCounter - config.mutex.Unlock() + return binaryData, rowCounter +} + +//export FLBPluginFlushCtx +func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int { + id := getFLBPluginContext(ctx) + config, ok := configMap[id] + if !ok { + log.Printf("Finding configuration for output instance with id: %d failed in FLBPluginFlushCtx", id) + return output.FLB_ERROR + } + + flushCtx, flushCancel := context.WithTimeout(context.Background(), config.flushTimeout) + defer flushCancel() + + // Drain ready responses and check whether a new stream should be created. + checkAllStreamResponses(flushCtx, &config.managedStreamSlice, false, &config.mutex, config.exactlyOnce, id) + createNewStreamDynamicScaling(flushCtx, config) + + // Reuse binaryData slice across flushes to avoid per-flush allocation. + // Nil out retained elements before reset so GC can reclaim previous batch's bytes. + for i := range config.binaryData { + config.binaryData[i] = nil } + binaryData := config.binaryData[:0] + + // Decode and serialize all records; mid-size chunks are sent inside the helper. + binaryData, rowCounter := decodeAndSerializeRecords(flushCtx, config, id, newDecoder(data, int(length)), binaryData) + + // Send the final (possibly partial) chunk. + flushChunk(flushCtx, config, id, binaryData, rowCounter) // Write back the grown slice so its capacity is reused on the next flush. config.binaryData = binaryData @@ -889,7 +882,7 @@ func FLBPluginExitCtx(ctx unsafe.Pointer) int { // Drain all pending responses before closing streams to avoid data loss. // waitForResponse=true blocks until every in-flight AppendRows result is received. checkAllStreamResponses(ms_ctx, &config.managedStreamSlice, true, &config.mutex, config.exactlyOnce, id) - errFlag := finalizeCloseAllStreams(&config, id) + errFlag := finalizeCloseAllStreams(config, id) if config.client != nil { if err := config.client.Close(); err != nil { diff --git a/out_writeapi_test.go b/out_writeapi_test.go index ace2af7..683f756 100644 --- a/out_writeapi_test.go +++ b/out_writeapi_test.go @@ -664,26 +664,27 @@ func TestFLBPluginFlushCtxDynamicScaling(t *testing.T) { defer patchRecord.Reset() // Creates new stream with a mock buildStream function to count the number of times we scale up - patchBuild := gomonkey.ApplyFunc(buildStream, func(ctx context.Context, config **outputConfig, streamIndex int) error { + origBuildStream := buildStream + buildStream = func(ctx context.Context, config *outputConfig, streamIndex int) error { checks.buildStreamCalled++ - currManagedStream, err := getWriter((*config).client, ctx, (*config).currProjectID, - managedwriter.WithType((*config).streamType), - managedwriter.WithDestinationTable((*config).tableRef), + currManagedStream, err := getWriter(config.client, ctx, config.currProjectID, + managedwriter.WithType(config.streamType), + managedwriter.WithDestinationTable(config.tableRef), // Use the descriptor proto when creating the new managed stream - managedwriter.WithSchemaDescriptor((*config).schemaDesc), - managedwriter.EnableWriteRetries((*config).enableRetry), - managedwriter.WithMaxInflightBytes((*config).maxQueueBytes), - managedwriter.WithMaxInflightRequests((*config).maxQueueRequests), + managedwriter.WithSchemaDescriptor(config.schemaDesc), + managedwriter.EnableWriteRetries(config.enableRetry), + managedwriter.WithMaxInflightBytes(config.maxQueueBytes), + managedwriter.WithMaxInflightRequests(config.maxQueueRequests), ) - streamSlice := *(*config).managedStreamSlice + streamSlice := *config.managedStreamSlice if err == nil { (streamSlice)[streamIndex].managedstream = currManagedStream } return nil - }) - defer patchBuild.Reset() + } + defer func() { buildStream = origBuildStream }() // Converts id (int) to type unsafe.Pointer to be used as the ctx // Use the address of setID instead of its value