Prevent hangs in OncePer when the callback panics

If the callback passed to Once panics it was leaving the waiter
in place that would never be completed.  Move writing the value
and signalling the waiter to defer.

Test: TestOncePerPanic
Change-Id: Icc4d3b779a79914fcd881d61d38dffcc2f591c39

2 files changed, 51 insertions, 5 deletions

diff --git a/android/onceper.go b/android/onceper.go
index 5ad17fa91..ff865c2e6 100644
--- a/android/onceper.go
+++ b/android/onceper.go
@@ -40,7 +40,8 @@ func (once *OncePer) maybeWaitFor(key OnceKey, value interface{}) interface{} {
 }
 
 // Once computes a value the first time it is called with a given key per OncePer, and returns the
-// value without recomputing when called with the same key.  key must be hashable.
+// value without recomputing when called with the same key.  key must be hashable.  If value panics
+// the panic will be propagated but the next call to Once with the same key will return nil.
 func (once *OncePer) Once(key OnceKey, value func() interface{}) interface{} {
 	// Fast path: check if the key is already in the map
 	if v, ok := once.values.Load(key); ok {
@@ -54,10 +55,15 @@ func (once *OncePer) Once(key OnceKey, value func() interface{}) interface{} {
 		return once.maybeWaitFor(key, v)
 	}
 
-	// The waiter is inserted, call the value constructor, store it, and signal the waiter
-	v := value()
-	once.values.Store(key, v)
-	close(waiter)
+	// The waiter is inserted, call the value constructor, store it, and signal the waiter.  Use defer in case
+	// the function panics.
+	var v interface{}
+	defer func() {
+		once.values.Store(key, v)
+		close(waiter)
+	}()
+
+	v = value()
 
 	return v
 }
diff --git a/android/onceper_test.go b/android/onceper_test.go
index 95303bafd..1a55ff4b1 100644
--- a/android/onceper_test.go
+++ b/android/onceper_test.go
@@ -175,3 +175,43 @@ func TestOncePerReentrant(t *testing.T) {
 		t.Errorf(`reentrant Once should return "a": %q`, a)
 	}
 }
+
+// Test that a recovered panic in a Once function doesn't deadlock
+func TestOncePerPanic(t *testing.T) {
+	once := OncePer{}
+	key := NewOnceKey("key")
+
+	ch := make(chan interface{})
+
+	var a interface{}
+
+	go func() {
+		defer func() {
+			ch <- recover()
+		}()
+
+		a = once.Once(key, func() interface{} {
+			panic("foo")
+		})
+	}()
+
+	p := <-ch
+
+	if p.(string) != "foo" {
+		t.Errorf(`expected panic with "foo", got %#v`, p)
+	}
+
+	if a != nil {
+		t.Errorf(`expected a to be nil, got %#v`, a)
+	}
+
+	// If the call to Once that panicked leaves the key in a bad state this will deadlock
+	b := once.Once(key, func() interface{} {
+		return "bar"
+	})
+
+	// The second call to Once should return nil inserted by the first call that panicked.
+	if b != nil {
+		t.Errorf(`expected b to be nil, got %#v`, b)
+	}
+}