From 28856d481e55547cc59fa038bf0e061308754fd7 Mon Sep 17 00:00:00 2001
From: Abhinav Gupta <abg@uber.com>
Date: Wed, 13 Jul 2016 13:40:13 -0700
Subject: [PATCH] transport/options: Make Options immutable

---
 transport/options.go      | 24 ++++++++-------
 transport/options_test.go | 65 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 11 deletions(-)
 create mode 100644 transport/options_test.go

diff --git a/transport/options.go b/transport/options.go
index 044c536a0..621babd0d 100644
--- a/transport/options.go
+++ b/transport/options.go
@@ -49,25 +49,27 @@ type Options struct {
 	data map[interface{}]interface{}
 }
 
-// With adds the given key-value pair to the option and returns the new
-// Options object.
+// With returns a copy of this Options object with the given key-value pair
+// added to it.
 //
 // The key should be a custom type to avoid conflicts with options of other
 // components.
 //
-// The returned object MAY not point to the same underlying data store as the
-// original Options so the returned Options MUST always be used instead of the
-// original.
-//
 // 	opts = opts.With(foo{}, bar)
 // 	opts = opts.With(baz{}, qux)
 //
-func (o Options) With(k, v interface{}) Options {
-	if o.data == nil {
-		o.data = make(map[interface{}]interface{})
+func (o Options) With(key, val interface{}) Options {
+	var data map[interface{}]interface{}
+	if o.data != nil {
+		data = make(map[interface{}]interface{}, len(o.data)+1)
+		for k, v := range o.data {
+			data[k] = v
+		}
+	} else {
+		data = make(map[interface{}]interface{})
 	}
-	o.data[k] = v
-	return o
+	data[key] = val
+	return Options{data}
 }
 
 // Get returns the value associated with the given key.
diff --git a/transport/options_test.go b/transport/options_test.go
new file mode 100644
index 000000000..f9b8caaaf
--- /dev/null
+++ b/transport/options_test.go
@@ -0,0 +1,65 @@
+// Copyright (c) 2016 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package transport
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestOptionsIsImmutable(t *testing.T) {
+	type foo struct{}
+	type bar struct{}
+
+	var o Options
+	withFoo := o.With(foo{}, "foo")
+	withBar := o.With(bar{}, "bar")
+	withFooBar := withFoo.With(bar{}, "foobar")
+
+	_, ok := o.Get(foo{})
+	assert.False(t, ok, "did not expect to find foo{} in o")
+
+	_, ok = o.Get(bar{})
+	assert.False(t, ok, "did not expect to find bar{} in o")
+
+	_, ok = withFoo.Get(bar{})
+	assert.False(t, ok, "did not expect to find bar{} in withFoo")
+
+	_, ok = withBar.Get(foo{})
+	assert.False(t, ok, "did not expect to find foo{} in withBar")
+
+	if v, ok := withFoo.Get(foo{}); assert.True(t, ok, "expected foo{} in withFoo") {
+		assert.Equal(t, "foo", v, "withFoo[foo{}] did not match")
+	}
+
+	if v, ok := withBar.Get(bar{}); assert.True(t, ok, "expected bar{} in withBar") {
+		assert.Equal(t, "bar", v, "withBar[bar{}] did not match")
+	}
+
+	if v, ok := withFooBar.Get(foo{}); assert.True(t, ok, "expected foo{} in withFooBar") {
+		assert.Equal(t, "foo", v, "withFooBar[foo{}] did not match")
+	}
+
+	if v, ok := withFooBar.Get(bar{}); assert.True(t, ok, "expected bar{} in withFooBar") {
+		assert.Equal(t, "foobar", v, "withFooBar[bar{}] did not match")
+	}
+}