internal/impl: avoid redundant lazy extension inits
After taking the lock on a lazy extension's state, check to see if it
was initialized while we were waiting for the lock.
Change-Id: I1cbd52e9d655eec6c9142c97689ae36f219a28f2
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/216898
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
diff --git a/internal/impl/codec_extension.go b/internal/impl/codec_extension.go
index 141a3cd..1098159 100644
--- a/internal/impl/codec_extension.go
+++ b/internal/impl/codec_extension.go
@@ -110,6 +110,9 @@
func (f *ExtensionField) lazyInit() {
f.lazy.mu.Lock()
defer f.lazy.mu.Unlock()
+ if atomic.LoadUint32(&f.lazy.atomicOnce) == 1 {
+ return
+ }
if f.lazy.xi != nil {
b := f.lazy.b
val := f.typ.New()
@@ -133,7 +136,7 @@
wtyp := wire.Type(tag & 7)
var out unmarshalOutput
var err error
- val, out, err = f.lazy.xi.funcs.unmarshal(b, val, num, wtyp, unmarshalOptions{}) // TODO: options
+ val, out, err = f.lazy.xi.funcs.unmarshal(b, val, num, wtyp, unmarshalOptions{})
if err != nil {
panic(errors.New("decode failure in lazy extension decoding: %v", err))
}
diff --git a/proto/extension_test.go b/proto/extension_test.go
index cf1ef6c..b3a0cff 100644
--- a/proto/extension_test.go
+++ b/proto/extension_test.go
@@ -6,6 +6,7 @@
import (
"fmt"
+ "sync"
"testing"
"github.com/google/go-cmp/cmp"
@@ -66,3 +67,39 @@
}
}
+
+func TestExtensionGetRace(t *testing.T) {
+ // Concurrently fetch an extension value while marshaling the message containing it.
+ // Create the message with proto.Unmarshal to give lazy extension decoding (if present)
+ // a chance to occur.
+ want := int32(42)
+ m1 := &testpb.TestAllExtensions{}
+ proto.SetExtension(m1, testpb.E_OptionalNestedMessageExtension, &testpb.TestAllExtensions_NestedMessage{A: proto.Int32(want)})
+ b, err := proto.Marshal(m1)
+ if err != nil {
+ t.Fatal(err)
+ }
+ m := &testpb.TestAllExtensions{}
+ if err := proto.Unmarshal(b, m); err != nil {
+ t.Fatal(err)
+ }
+ var wg sync.WaitGroup
+ for i := 0; i < 3; i++ {
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ if _, err := proto.Marshal(m); err != nil {
+ t.Error(err)
+ }
+ }()
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ got := proto.GetExtension(m, testpb.E_OptionalNestedMessageExtension).(*testpb.TestAllExtensions_NestedMessage).GetA()
+ if got != want {
+ t.Errorf("GetExtension(optional_nested_message).a = %v, want %v", got, want)
+ }
+ }()
+ }
+ wg.Wait()
+}
