diff --git a/map_to_proto.go b/map_to_proto.go index 72f1f59..35c501d 100644 --- a/map_to_proto.go +++ b/map_to_proto.go @@ -89,6 +89,29 @@ func getPooledMessage(md protoreflect.MessageDescriptor) *dynamicpb.Message { return p.(*sync.Pool).Get().(*dynamicpb.Message) } +// cleanupMsgPool removes all messagePool entries for the descriptor tree +// rooted at md. Call this when an outputConfig is torn down to prevent +// pool entries (and their associated descriptor references) from accumulating +// across init/exit cycles. +func cleanupMsgPool(md protoreflect.MessageDescriptor) { + cleanupMsgPoolRecursive(md, make(map[protoreflect.MessageDescriptor]bool)) +} + +func cleanupMsgPoolRecursive(md protoreflect.MessageDescriptor, visited map[protoreflect.MessageDescriptor]bool) { + if visited[md] { + return + } + visited[md] = true + messagePool.Delete(md) + fields := md.Fields() + for i := 0; i < fields.Len(); i++ { + f := fields.Get(i) + if f.Kind() == protoreflect.MessageKind || f.Kind() == protoreflect.GroupKind { + cleanupMsgPoolRecursive(f.Message(), visited) + } + } +} + // 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. diff --git a/map_to_proto_test.go b/map_to_proto_test.go index 4fbfb37..f6c292d 100644 --- a/map_to_proto_test.go +++ b/map_to_proto_test.go @@ -1141,6 +1141,68 @@ func TestMessagePoolReuse_RawPath_NestedRecord(t *testing.T) { assert.Equal(t, int64(0), msg2.Get(md.Fields().ByName("nested")).Message().Get(nestedMd.Fields().ByName("y")).Int()) } +// TestCleanupMsgPool verifies that cleanupMsgPool removes all messagePool +// entries for a flat schema and a nested (STRUCT) schema. +func TestCleanupMsgPool(t *testing.T) { + t.Run("flat schema", func(t *testing.T) { + schema := &storagepb.TableSchema{ + Fields: []*storagepb.TableFieldSchema{ + {Name: "name", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE}, + {Name: "count", Type: storagepb.TableFieldSchema_INT64, Mode: storagepb.TableFieldSchema_NULLABLE}, + }, + } + md := buildMD(t, schema) + cache := buildFieldLookupCache(md) + + // Populate the pool by serialising one record. + _, err := mapToBinary(md, map[string]interface{}{"name": "x", "count": int64(1)}, cache) + require.NoError(t, err) + + _, populated := messagePool.Load(md) + assert.True(t, populated, "messagePool should have entry before cleanup") + + cleanupMsgPool(md) + + _, populated = messagePool.Load(md) + assert.False(t, populated, "messagePool should have no entry after cleanup") + }) + + t.Run("nested schema cleans sub-message entries", func(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}, + }, + }, + }, + } + md := buildMD(t, schema) + nestedMd := md.Fields().ByName("nested").Message() + cache := buildFieldLookupCache(md) + + _, err := rawMapToBinary(md, map[interface{}]interface{}{ + "nested": map[interface{}]interface{}{"x": []byte("hello")}, + }, cache) + require.NoError(t, err) + + _, topPopulated := messagePool.Load(md) + _, subPopulated := messagePool.Load(nestedMd) + assert.True(t, topPopulated, "top-level entry should exist before cleanup") + assert.True(t, subPopulated, "sub-message entry should exist before cleanup") + + cleanupMsgPool(md) + + _, topPopulated = messagePool.Load(md) + _, subPopulated = messagePool.Load(nestedMd) + assert.False(t, topPopulated, "top-level entry should be removed after cleanup") + assert.False(t, subPopulated, "sub-message entry should be removed after cleanup") + }) +} + // 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 b0a3902..4d3df50 100644 --- a/out_writeapi.go +++ b/out_writeapi.go @@ -491,8 +491,14 @@ func getLeastLoadedStream(streamSlice *[]*streamConfig) int { // This is a test-only method which takes in a config id and returns the current offset value of the struct corresponding to the id func getOffset(id int) int64 { - config := configMap[id] + config, ok := configMap[id] + if !ok || config == nil { + return 0 + } streamSlice := *config.managedStreamSlice + if len(streamSlice) == 0 { + return 0 + } return streamSlice[0].offsetCounter } @@ -602,7 +608,12 @@ var newDecoder = func(data unsafe.Pointer, length int) *output.FLBDecoder { // Mock it whenever needed var getFLBPluginContext = func(ctx unsafe.Pointer) int { if ctx != nil { - return output.FLBPluginGetContext(ctx).(int) + id, ok := output.FLBPluginGetContext(ctx).(int) + if !ok { + log.Printf("FLBPluginGetContext returned unexpected type for context pointer %p", ctx) + return 0 + } + return id } return 0 } @@ -895,6 +906,11 @@ func FLBPluginExitCtx(ctx unsafe.Pointer) int { if errFlag { return output.FLB_ERROR } + + // Release messagePool entries for this config's descriptor tree so that + // the descriptors and their cached *dynamicpb.Message objects can be GC'd. + cleanupMsgPool(config.messageDescriptor) + return output.FLB_OK }