From 0467ab0356363e37aff0f1bfa7f40e06e1362b2b Mon Sep 17 00:00:00 2001
From: "zhongyang.wu" <zhongyang.wu@outlook.com>
Date: Mon, 7 Jun 2021 21:24:11 -0400
Subject: [PATCH] fix: TraceState cannot insert new key-value pairs.

---
 opentelemetry/src/sdk/trace/tracer.rs   |  2 +-
 opentelemetry/src/trace/span_context.rs | 23 +++++++++++++++++++----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/opentelemetry/src/sdk/trace/tracer.rs b/opentelemetry/src/sdk/trace/tracer.rs
index b74d8cb4f0..46d401db7a 100644
--- a/opentelemetry/src/sdk/trace/tracer.rs
+++ b/opentelemetry/src/sdk/trace/tracer.rs
@@ -361,7 +361,7 @@ mod tests {
             SamplingResult {
                 decision: SamplingDecision::RecordAndSample,
                 attributes: Vec::new(),
-                trace_state: trace_state.insert("foo".into(), "notbar".into()).unwrap(),
+                trace_state: trace_state.insert("foo", "notbar").unwrap(),
             }
         }
     }
diff --git a/opentelemetry/src/trace/span_context.rs b/opentelemetry/src/trace/span_context.rs
index c226e12536..4ab7c5f9f6 100644
--- a/opentelemetry/src/trace/span_context.rs
+++ b/opentelemetry/src/trace/span_context.rs
@@ -293,15 +293,16 @@ impl TraceState {
     ///
     /// ['spec']: https://www.w3.org/TR/trace-context/#list
     #[allow(clippy::all)]
-    pub fn insert(&self, key: String, value: String) -> Result<TraceState, ()> {
+    pub fn insert<T: Into<String>>(&self, key: T, value: T) -> Result<TraceState, ()> {
+        let (key, value) = (key.into(), value.into());
         if !TraceState::valid_key(key.as_str()) || !TraceState::valid_value(value.as_str()) {
             return Err(());
         }
 
-        let mut trace_state = self.delete(key.clone())?;
+        let mut trace_state = self.delete_from_deque(key.clone())?;
         let kvs = trace_state.0.get_or_insert(VecDeque::with_capacity(1));
 
-        kvs.push_front((key, value));
+        kvs.push_front((key.into(), value.into()));
 
         Ok(trace_state)
     }
@@ -312,11 +313,17 @@ impl TraceState {
     ///
     /// ['spec']: https://www.w3.org/TR/trace-context/#list
     #[allow(clippy::all)]
-    pub fn delete(&self, key: String) -> Result<TraceState, ()> {
+    pub fn delete<T: Into<String>>(&self, key: T) -> Result<TraceState, ()> {
+        let key = key.into();
         if !TraceState::valid_key(key.as_str()) {
             return Err(());
         }
 
+        self.delete_from_deque(key)
+    }
+
+    /// Delete key from trace state's deque. The key MUST be valid
+    fn delete_from_deque(&self, key: String) -> Result<TraceState, ()> {
         let mut owned = self.clone();
         let kvs = owned.0.as_mut().ok_or(())?;
 
@@ -533,4 +540,12 @@ mod tests {
             assert!(deleted_trace_state.get(test_case.2).is_none());
         }
     }
+
+    #[test]
+    fn test_trace_state_insert() {
+        let trace_state = TraceState::from_key_value(vec![("foo", "bar")]).unwrap();
+        let inserted_trace_state = trace_state.insert("testkey", "testvalue").unwrap();
+        assert!(trace_state.get("testkey").is_none()); // The original state doesn't change
+        assert_eq!(inserted_trace_state.get("testkey").unwrap(), "testvalue"); //
+    }
 }