From 136ac4728de82531e2642c8ae75a4c6c248f2c62 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 4 Oct 2016 20:20:07 -0700
Subject: [PATCH] helper/shadow: KeyedValue.Close

---
 helper/shadow/keyed_value.go      | 23 +++++++++++
 helper/shadow/keyed_value_test.go | 65 +++++++++++++++++++++++++++++++
 helper/shadow/value_test.go       | 31 +++++++++++++++
 3 files changed, 119 insertions(+)

diff --git a/helper/shadow/keyed_value.go b/helper/shadow/keyed_value.go
index 9c6d576e9..c0f6b58f0 100644
--- a/helper/shadow/keyed_value.go
+++ b/helper/shadow/keyed_value.go
@@ -11,6 +11,24 @@ type KeyedValue struct {
 	once    sync.Once
 	values  map[string]interface{}
 	waiters map[string]*Value
+	closed  bool
+}
+
+// Close closes the value. This can never fail. For a definition of
+// "close" see the ErrClosed docs.
+func (w *KeyedValue) Close() error {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	// Set closed to true always
+	w.closed = true
+
+	// For all waiters, complete with ErrClosed
+	for _, w := range w.waiters {
+		w.SetValue(ErrClosed)
+	}
+
+	return nil
 }
 
 // Value returns the value that was set for the given key, or blocks
@@ -62,6 +80,11 @@ func (w *KeyedValue) valueWaiter(k string) (interface{}, *Value) {
 		return v, nil
 	}
 
+	// If we're closed, return that
+	if w.closed {
+		return ErrClosed, nil
+	}
+
 	// No pending value, check for a waiter
 	val := w.waiters[k]
 	if val == nil {
diff --git a/helper/shadow/keyed_value_test.go b/helper/shadow/keyed_value_test.go
index ff59fe8de..a141606a2 100644
--- a/helper/shadow/keyed_value_test.go
+++ b/helper/shadow/keyed_value_test.go
@@ -73,3 +73,68 @@ func TestKeyedValueOk(t *testing.T) {
 		t.Fatalf("bad: %#v", val)
 	}
 }
+
+func TestKeyedValueClose(t *testing.T) {
+	var v KeyedValue
+
+	// Close
+	v.Close()
+
+	// Try again
+	val, ok := v.ValueOk("foo")
+	if !ok {
+		t.Fatal("should be ok")
+	}
+
+	// Verify
+	if val != ErrClosed {
+		t.Fatalf("bad: %#v", val)
+	}
+}
+
+func TestKeyedValueClose_blocked(t *testing.T) {
+	var v KeyedValue
+
+	// Start reading this should be blocking
+	valueCh := make(chan interface{})
+	go func() {
+		valueCh <- v.Value("foo")
+	}()
+
+	// We should not get the value
+	select {
+	case <-valueCh:
+		t.Fatal("shouldn't receive value")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	// Close
+	v.Close()
+
+	// Verify
+	val := <-valueCh
+	if val != ErrClosed {
+		t.Fatalf("bad: %#v", val)
+	}
+}
+
+func TestKeyedValueClose_existing(t *testing.T) {
+	var v KeyedValue
+
+	// Set a value
+	v.SetValue("foo", "bar")
+
+	// Close
+	v.Close()
+
+	// Try again
+	val, ok := v.ValueOk("foo")
+	if !ok {
+		t.Fatal("should be ok")
+	}
+
+	// Verify
+	if val != "bar" {
+		t.Fatalf("bad: %#v", val)
+	}
+}
diff --git a/helper/shadow/value_test.go b/helper/shadow/value_test.go
index ad8a64aa9..8bc957939 100644
--- a/helper/shadow/value_test.go
+++ b/helper/shadow/value_test.go
@@ -55,6 +55,37 @@ func TestValueClose(t *testing.T) {
 	}
 }
 
+func TestValueClose_blocked(t *testing.T) {
+	var v Value
+
+	// Start trying to get the value
+	valueCh := make(chan interface{})
+	go func() {
+		valueCh <- v.Value()
+	}()
+
+	// We should not get the value
+	select {
+	case <-valueCh:
+		t.Fatal("shouldn't receive value")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	// Set the value
+	v.Close()
+	val := <-valueCh
+
+	// Verify
+	if val != ErrClosed {
+		t.Fatalf("bad: %#v", val)
+	}
+
+	// We should be able to ask for the value again immediately
+	if val := v.Value(); val != ErrClosed {
+		t.Fatalf("bad: %#v", val)
+	}
+}
+
 func TestValueClose_existing(t *testing.T) {
 	var v Value