From 89e13a1d94c230cd1552d5a73e6730dace710921 Mon Sep 17 00:00:00 2001
From: Tigran Najaryan <4194920+tigrannajaryan@users.noreply.github.com>
Date: Fri, 7 Jul 2023 14:32:52 -0400
Subject: [PATCH] [pdata] Fix Map.EnsureCapacity bug when the Map is not empty
 (#8040)

EnsureCapacity was incorrectly clearing the Map when increasing the
capacity. It will now correctly preserve existing elements.

---------

Co-authored-by: Alex Boten <alex@boten.ca>
---
 pdata/pcommon/map.go      |  6 +++---
 pdata/pcommon/map_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/pdata/pcommon/map.go b/pdata/pcommon/map.go
index 97cc505f7dc..8e1b2bbccb6 100644
--- a/pdata/pcommon/map.go
+++ b/pdata/pcommon/map.go
@@ -35,11 +35,11 @@ func (m Map) Clear() {
 // EnsureCapacity increases the capacity of this Map instance, if necessary,
 // to ensure that it can hold at least the number of elements specified by the capacity argument.
 func (m Map) EnsureCapacity(capacity int) {
-	if capacity <= cap(*m.getOrig()) {
+	oldOrig := *m.getOrig()
+	if capacity <= cap(oldOrig) {
 		return
 	}
-	oldOrig := *m.getOrig()
-	*m.getOrig() = make([]otlpcommon.KeyValue, 0, capacity)
+	*m.getOrig() = make([]otlpcommon.KeyValue, len(oldOrig), capacity)
 	copy(*m.getOrig(), oldOrig)
 }
 
diff --git a/pdata/pcommon/map_test.go b/pdata/pcommon/map_test.go
index 3201d4f0ea0..28f5a600906 100644
--- a/pdata/pcommon/map_test.go
+++ b/pdata/pcommon/map_test.go
@@ -375,6 +375,38 @@ func TestMap_EnsureCapacity(t *testing.T) {
 	assert.Equal(t, 8, cap(*am.getOrig()))
 }
 
+func TestMap_EnsureCapacity_Existing(t *testing.T) {
+	am := NewMap()
+	am.PutStr("foo", "bar")
+
+	assert.Equal(t, 1, am.Len())
+
+	// Add more capacity.
+	am.EnsureCapacity(5)
+
+	// Ensure previously existing element is still there.
+	assert.Equal(t, 1, am.Len())
+	v, ok := am.Get("foo")
+	assert.Equal(t, v.Str(), "bar")
+	assert.True(t, ok)
+
+	assert.Equal(t, 5, cap(*am.getOrig()))
+
+	// Add one more element.
+	am.PutStr("abc", "xyz")
+
+	// Verify that both elements are there.
+	assert.Equal(t, 2, am.Len())
+
+	v, ok = am.Get("foo")
+	assert.Equal(t, v.Str(), "bar")
+	assert.True(t, ok)
+
+	v, ok = am.Get("abc")
+	assert.Equal(t, v.Str(), "xyz")
+	assert.True(t, ok)
+}
+
 func TestMap_Clear(t *testing.T) {
 	am := NewMap()
 	assert.Nil(t, *am.getOrig())