From 7c867de9fd6c4da147166e6c68ceecf4b99f139a Mon Sep 17 00:00:00 2001
From: zacscoding <zaccoding725@gmail.com>
Date: Mon, 29 Mar 2021 20:49:13 +0900
Subject: [PATCH] Add delete a baggage item in span

Signed-off-by: zacscoding <zaccoding725@gmail.com>
---
 span_context.go      | 13 +++++++++++++
 span_context_test.go | 13 +++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/span_context.go b/span_context.go
index d3627624..5b2307be 100644
--- a/span_context.go
+++ b/span_context.go
@@ -314,6 +314,7 @@ func (c *SpanContext) CopyFrom(ctx *SpanContext) {
 }
 
 // WithBaggageItem creates a new context with an extra baggage item.
+// Delete a baggage item if provided blank value.
 //
 // The SpanContext is designed to be immutable and passed by value. As such,
 // it cannot contain any locks, and should only hold immutable data, including baggage.
@@ -324,6 +325,18 @@ func (c *SpanContext) CopyFrom(ctx *SpanContext) {
 // a trace), or it needs to do a copy-on-write, which is the approach taken here.
 func (c SpanContext) WithBaggageItem(key, value string) SpanContext {
 	var newBaggage map[string]string
+	// unset baggage item
+	if value == "" {
+		if _, ok := c.baggage[key]; !ok {
+			return c
+		}
+		newBaggage = make(map[string]string, len(c.baggage))
+		for k, v := range c.baggage {
+			newBaggage[k] = v
+		}
+		delete(newBaggage, key)
+		return SpanContext{c.traceID, c.spanID, c.parentID, newBaggage, "", c.samplingState, c.remote}
+	}
 	if c.baggage == nil {
 		newBaggage = map[string]string{key: value}
 	} else {
diff --git a/span_context_test.go b/span_context_test.go
index cf6584b7..14213e42 100644
--- a/span_context_test.go
+++ b/span_context_test.go
@@ -79,6 +79,19 @@ func TestSpanContext_WithBaggageItem(t *testing.T) {
 	assert.Equal(t, map[string]string{"some-KEY": "Some-Other-Value"}, ctx.baggage)
 }
 
+func TestSpanContext_WithBaggageItem_Delete(t *testing.T) {
+	var ctx SpanContext
+	ctx = ctx.WithBaggageItem("some-KEY", "")
+	assert.Nil(t, ctx.baggage)
+	ctx = ctx.WithBaggageItem("some-KEY", "Some-Value")
+	assert.Equal(t, map[string]string{"some-KEY": "Some-Value"}, ctx.baggage)
+	ctx = ctx.WithBaggageItem("another-KEY", "")
+	assert.Equal(t, map[string]string{"some-KEY": "Some-Value"}, ctx.baggage)
+	ctx2 := ctx.WithBaggageItem("some-KEY", "")
+	assert.Equal(t, map[string]string{"some-KEY": "Some-Value"}, ctx.baggage, "parent unchanged")
+	assert.Equal(t, map[string]string{}, ctx2.baggage)
+}
+
 func TestSpanContext_Flags(t *testing.T) {
 
 	var tests = map[string]struct {