internal/legacy: guarantee one-to-one generation of wrapper types

When concurrent requests to wrap a type occurs, it is possible that
two different descriptors are created for the same Go type.
Add sufficient synchronization to ensure that one descriptor is ever
returned for one Go type.

Change-Id: Idbbca4c1877a70317b39900ae83bfc3085d4a9c5
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/168398
Reviewed-by: Herbie Ong <herbie@google.com>
diff --git a/internal/legacy/extension.go b/internal/legacy/extension.go
index 839d597..0c88c83 100644
--- a/internal/legacy/extension.go
+++ b/internal/legacy/extension.go
@@ -134,7 +134,9 @@
 		Tag:           ptag.Marshal(t, enumName),
 		Filename:      filename,
 	}
-	extensionDescCache.Store(t, d)
+	if d, ok := extensionDescCache.LoadOrStore(t, d); ok {
+		return d.(*papi.ExtensionDesc)
+	}
 	return d
 }
 
@@ -145,10 +147,6 @@
 func extensionTypeFromDesc(d *papi.ExtensionDesc) pref.ExtensionType {
 	// Fast-path: check whether an extension type is already nested within.
 	if d.Type != nil {
-		// Cache descriptor for future extensionDescFromType operation.
-		// This assumes that there is only one legacy protoapi.ExtensionDesc
-		// that wraps any given specific protoreflect.ExtensionType.
-		extensionDescCache.LoadOrStore(d.Type, d)
 		return d.Type
 	}
 
@@ -192,8 +190,10 @@
 	xt := pimpl.Export{}.ExtensionTypeOf(xd, zv)
 
 	// Cache the conversion for both directions.
-	extensionDescCache.Store(xt, d)
-	extensionTypeCache.Store(dk, xt)
+	extensionDescCache.LoadOrStore(xt, d)
+	if xt, ok := extensionTypeCache.LoadOrStore(dk, xt); ok {
+		return xt.(pref.ExtensionType)
+	}
 	return xt
 }