diff --git a/map_to_proto.go b/map_to_proto.go index 9cb7826..72f1f59 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" @@ -91,8 +92,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,16 +173,38 @@ 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 } +// 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 { @@ -196,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) } @@ -247,26 +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 := dynamicpb.NewMessage(fd.Message()) - if err := rawPopulateMessage(subMsg, fd.Message(), sub, cache); err != nil { - 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 { @@ -315,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) } @@ -335,50 +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 := dynamicpb.NewMessage(fd.Message()) - if err := rawPopulateMessage(subMsg, fd.Message(), sub, cache); err != nil { - 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. @@ -414,7 +390,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: @@ -458,7 +434,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) } @@ -487,7 +472,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) } @@ -512,7 +501,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) } @@ -537,7 +530,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) } @@ -562,7 +559,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) } @@ -585,7 +586,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) } @@ -610,7 +615,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) } @@ -656,7 +665,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/map_to_proto_test.go b/map_to_proto_test.go index 583544d..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" @@ -880,7 +881,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. @@ -1009,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{ diff --git a/out_writeapi.go b/out_writeapi.go index 9c8474b..2712620 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -21,7 +21,9 @@ import ( "log" "math" "strconv" + "strings" "sync" + "sync/atomic" "time" "unsafe" @@ -62,25 +64,28 @@ type outputConfig struct { numRetries int timestampFields []string fieldCache fieldLookupCache + flushTimeout time.Duration + binaryData [][]byte // reused across flushes to avoid per-flush allocation } var ( ms_ctx = context.Background() configMap = make(map[int]*outputConfig) - configID = 0 + configID atomic.Int64 ) 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 + 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. @@ -187,12 +192,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) + } } } } @@ -271,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 @@ -327,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 @@ -361,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 { @@ -380,16 +401,16 @@ 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 { break } - // Unsuccesful data append + // Unsuccessful data append if rebuildPredicate(err) { currStream.managedstream.Finalize(ctx) currStream.managedstream.Close() @@ -414,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 { @@ -429,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) @@ -447,11 +468,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 } } @@ -476,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(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{ @@ -490,14 +511,14 @@ func createNewStreamDynamicScaling(config **outputConfig) { appendResults: &newResQueue, } - if mostEfficientQueueLength > (*config).requestCountThreshold { - *(*config).managedStreamSlice = append(*(*config).managedStreamSlice, &newStream) - newStreamIndex := len(*(*config).managedStreamSlice) - 1 - err := buildStream(ms_ctx, config, newStreamIndex) + 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] } } @@ -562,6 +583,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 { @@ -572,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 @@ -652,6 +678,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 { @@ -708,23 +741,21 @@ func FLBPluginInit(plugin unsafe.Pointer) int { managedStreamSlice: &streamSlice, timestampFields: timestampFields, fieldCache: buildFieldLookupCache(md), + flushTimeout: time.Duration(flushTimeoutSec) * time.Second, } // 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 } - 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 } @@ -735,86 +766,98 @@ 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 - } +// 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() - // Calls checkResponses for all streams in slice - checkAllStreamResponses(ms_ctx, &config.managedStreamSlice, false, &config.mutex, config.exactlyOnce, id) - // Checks for need to dynamically scale - createNewStreamDynamicScaling(&config) + 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 + } + if config.exactlyOnce { + config.mutex.Lock() + (*config.managedStreamSlice)[streamIndex].offsetCounter += rowCount + config.mutex.Unlock() + } +} - // Create Fluent Bit decoder - dec := output.NewDecoder(data, int(length)) - // Pre-allocate binaryData slice to reduce append-driven growth (#5) - binaryData := make([][]byte, 0, 256) +// 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 - // Find stream with least number of awaiting queue responses - config.mutex.Lock() - leastLoadedStreamIndex := getLeastLoadedStream(config.managedStreamSlice) - config.mutex.Unlock() - - // 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 { - // Appending Rows - err := sendRequest(ms_ctx, 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 - - // Reuse the underlying array instead of nil to avoid re-allocation - 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++ } - // Appending Rows - err := sendRequest(ms_ctx, 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 + return output.FLB_OK } @@ -827,7 +870,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] @@ -836,10 +879,10 @@ 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) - errFlag := finalizeCloseAllStreams(&config, 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 { if err := config.client.Close(); err != nil { diff --git a/out_writeapi_test.go b/out_writeapi_test.go index f71ca5e..683f756 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() } @@ -446,11 +452,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 +638,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 @@ -656,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 @@ -814,11 +823,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 +1005,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 @@ -1046,6 +1057,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) { @@ -1115,6 +1240,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"},