From 04f281d0fccebb045f1cb25503e3d1cc8ebc8e2c Mon Sep 17 00:00:00 2001
From: FillZpp <FillZpp.pub@gmail.com>
Date: Mon, 9 May 2022 14:20:37 +0800
Subject: [PATCH] Fix panic for lazy dynamicRESTMapper

Signed-off-by: FillZpp <FillZpp.pub@gmail.com>
---
 pkg/client/apiutil/apiutil_suite_test.go     |  2 +-
 pkg/client/apiutil/dynamicrestmapper.go      | 19 +++++++++---
 pkg/client/apiutil/dynamicrestmapper_test.go | 32 +++++++++++++++-----
 3 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/pkg/client/apiutil/apiutil_suite_test.go b/pkg/client/apiutil/apiutil_suite_test.go
index 757bc42fec..f617195724 100644
--- a/pkg/client/apiutil/apiutil_suite_test.go
+++ b/pkg/client/apiutil/apiutil_suite_test.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package apiutil_test
+package apiutil
 
 import (
 	"testing"
diff --git a/pkg/client/apiutil/dynamicrestmapper.go b/pkg/client/apiutil/dynamicrestmapper.go
index 56a00371ff..e6cc51c6e9 100644
--- a/pkg/client/apiutil/dynamicrestmapper.go
+++ b/pkg/client/apiutil/dynamicrestmapper.go
@@ -19,6 +19,7 @@ package apiutil
 import (
 	"errors"
 	"sync"
+	"sync/atomic"
 
 	"golang.org/x/time/rate"
 	"k8s.io/apimachinery/pkg/api/meta"
@@ -38,7 +39,8 @@ type dynamicRESTMapper struct {
 
 	lazy bool
 	// Used for lazy init.
-	initOnce sync.Once
+	inited  uint32
+	initMtx sync.Mutex
 }
 
 // DynamicRESTMapperOption is a functional option on the dynamicRESTMapper.
@@ -125,11 +127,18 @@ func (drm *dynamicRESTMapper) setStaticMapper() error {
 
 // init initializes drm only once if drm is lazy.
 func (drm *dynamicRESTMapper) init() (err error) {
-	drm.initOnce.Do(func() {
-		if drm.lazy {
-			err = drm.setStaticMapper()
+	// skip init if drm is not lazy or has initialized
+	if !drm.lazy || atomic.LoadUint32(&drm.inited) != 0 {
+		return nil
+	}
+
+	drm.initMtx.Lock()
+	defer drm.initMtx.Unlock()
+	if drm.inited == 0 {
+		if err = drm.setStaticMapper(); err == nil {
+			atomic.StoreUint32(&drm.inited, 1)
 		}
-	})
+	}
 	return err
 }
 
diff --git a/pkg/client/apiutil/dynamicrestmapper_test.go b/pkg/client/apiutil/dynamicrestmapper_test.go
index 1e294b793d..6b88a3aa5f 100644
--- a/pkg/client/apiutil/dynamicrestmapper_test.go
+++ b/pkg/client/apiutil/dynamicrestmapper_test.go
@@ -14,9 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package apiutil_test
+package apiutil
 
 import (
+	"fmt"
 	"time"
 
 	. "github.com/onsi/ginkgo"
@@ -26,8 +27,6 @@ import (
 	"golang.org/x/time/rate"
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/runtime/schema"
-
-	"sigs.k8s.io/controller-runtime/pkg/client/apiutil"
 )
 
 var (
@@ -52,7 +51,7 @@ var _ = Describe("Dynamic REST Mapper", func() {
 		}
 
 		lim = rate.NewLimiter(rate.Limit(5), 5)
-		mapper, err = apiutil.NewDynamicRESTMapper(cfg, apiutil.WithLimiter(lim), apiutil.WithCustomMapper(func() (meta.RESTMapper, error) {
+		mapper, err = NewDynamicRESTMapper(cfg, WithLimiter(lim), WithCustomMapper(func() (meta.RESTMapper, error) {
 			baseMapper := meta.NewDefaultRESTMapper(nil)
 			addToMapper(baseMapper)
 
@@ -146,11 +145,28 @@ var _ = Describe("Dynamic REST Mapper", func() {
 			By("ensuring that it was only refreshed once")
 			Expect(count).To(Equal(1))
 		})
-	}
 
-	PIt("should lazily initialize if the lazy option is used", func() {
-
-	})
+		It("should lazily initialize if the lazy option is used", func() {
+			var err error
+			var failedOnce bool
+			mockErr := fmt.Errorf("mock failed once")
+			mapper, err = NewDynamicRESTMapper(cfg, WithLazyDiscovery, WithCustomMapper(func() (meta.RESTMapper, error) {
+				// Make newMapper fail once
+				if !failedOnce {
+					failedOnce = true
+					return nil, mockErr
+				}
+				baseMapper := meta.NewDefaultRESTMapper(nil)
+				addToMapper(baseMapper)
+				return baseMapper, nil
+			}))
+			Expect(err).NotTo(HaveOccurred())
+			Expect(mapper.(*dynamicRESTMapper).staticMapper).To(BeNil())
+
+			Expect(callWithTarget()).To(MatchError(mockErr))
+			Expect(callWithTarget()).To(Succeed())
+		})
+	}
 
 	Describe("KindFor", func() {
 		mapperTest(func() error {