Merge "Prevent hangs in OncePer when the callback panics"

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
+	// The waiter is inserted, call the value constructor, store it, and signal the waiter.  Use defer in case
-	v := value()
+	// the function panics.
-	once.values.Store(key, v)
+	var v interface{}
-	close(waiter)
+	defer func() {
+		once.values.Store(key, v)
+		close(waiter)
+	}()
+
+	v = value()
 
 	return v
 }

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)
+	}
+}