fix: guard type assertions and nil dereferences, release messagePool on exit

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.

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{

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
 }