Refactored code according to comments

Author: RanVaknin

core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java

@@ -28,7 +28,7 @@
 
 /**
  * The {@code TraceIdExecutionInterceptor} copies the trace details to the {@link #TRACE_ID_HEADER} header, assuming we seem to
- * be running in a lambda environment.
+ * be running in a lambda environment.`
  */
 @SdkProtectedApi
 public class TraceIdExecutionInterceptor implements ExecutionInterceptor {

utils-lite/src/main/java/software/amazon/awssdk/utilslite/SdkInternalThreadLocal.java

@@ -32,37 +32,50 @@ private SdkInternalThreadLocal() {
     }
 
     public static void put(String key, String value) {
-        Map<String, String> map = STORAGE.get();
-        if (map == null) {
-            map = new HashMap<>();
-            STORAGE.set(map);
+        if (key == null) {
+            throw new IllegalArgumentException("Key cannot be null");
         }
-
+        Map<String, String> map = STORAGE.get();
         if (value == null) {
-            map.remove(key);
+            if (map != null) {
+                map.remove(key);
+                if (map.isEmpty()) {
+                    STORAGE.remove();
+                }
+            }
         } else {
+            if (map == null) {
+                map = new HashMap<>();
+                STORAGE.set(map);
+            }
             map.put(key, value);
         }
     }
 
     public static String get(String key) {
+        if (key == null) {
+            throw new IllegalArgumentException("Key cannot be null");
+        }
         Map<String, String> map = STORAGE.get();
         return map != null ? map.get(key) : null;
     }
 
     public static String remove(String key) {
+        if (key == null) {
+            throw new IllegalArgumentException("Key cannot be null");
+        }
         Map<String, String> map = STORAGE.get();
         return map != null ? map.remove(key) : null;
     }
 
     public static void clear() {
-        Map<String, String> map = STORAGE.get();
-        if (map != null) {
-            map.clear();
-        }
+        STORAGE.remove();
     }
 
     public static boolean containsKey(String key) {
+        if (key == null) {
+            throw new IllegalArgumentException("Key cannot be null");
+        }
         Map<String, String> map = STORAGE.get();
         return map != null && map.containsKey(key);
     }