Merge pull request #31199 from TerryHowe/fix-flaky-registry-data-race

fix: flaky registry data race on mockdns close

pkg/registry/client_http_test.go

@@ -32,10 +32,7 @@ type HTTPRegistryClientTestSuite struct {
 
 func (suite *HTTPRegistryClientTestSuite) SetupSuite() {
 	// init test client
-	dockerRegistry := setup(&suite.TestSuite, false, false)
-
-	// Start Docker registry
-	go dockerRegistry.ListenAndServe()
+	setup(&suite.TestSuite, false, false)
 }
 
 func (suite *HTTPRegistryClientTestSuite) TearDownSuite() {

pkg/registry/client_insecure_tls_test.go

@@ -29,10 +29,7 @@ type InsecureTLSRegistryClientTestSuite struct {
 
 func (suite *InsecureTLSRegistryClientTestSuite) SetupSuite() {
 	// init test client
-	dockerRegistry := setup(&suite.TestSuite, true, true)
-
-	// Start Docker registry
-	go dockerRegistry.ListenAndServe()
+	setup(&suite.TestSuite, true, true)
 }
 
 func (suite *InsecureTLSRegistryClientTestSuite) TearDownSuite() {

pkg/registry/client_tls_test.go

@@ -31,10 +31,7 @@ type TLSRegistryClientTestSuite struct {
 
 func (suite *TLSRegistryClientTestSuite) SetupSuite() {
 	// init test client
-	dockerRegistry := setup(&suite.TestSuite, true, false)
-
-	// Start Docker registry
-	go dockerRegistry.ListenAndServe()
+	setup(&suite.TestSuite, true, false)
 }
 
 func (suite *TLSRegistryClientTestSuite) TearDownSuite() {

pkg/registry/main_test.go

@@ -0,0 +1,51 @@
+/*
+Copyright The Helm Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package registry
+
+import (
+	"net"
+	"os"
+	"testing"
+
+	"github.com/foxcpp/go-mockdns"
+)
+
+func TestMain(m *testing.M) {
+	// A mock DNS server needed for TLS connection testing.
+	var srv *mockdns.Server
+	var err error
+
+	srv, err = mockdns.NewServer(map[string]mockdns.Zone{
+		"helm-test-registry.": {
+			A: []string{"127.0.0.1"},
+		},
+	}, false)
+	if err != nil {
+		panic(err)
+	}
+
+	saveDialFunction := net.DefaultResolver.Dial
+	srv.PatchNet(net.DefaultResolver)
+
+	// Run all tests in the package
+	code := m.Run()
+
+	net.DefaultResolver.Dial = saveDialFunction
+	_ = srv.Close()
+
+	os.Exit(code)
+}

pkg/registry/utils_test.go

@@ -29,14 +29,12 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/distribution/distribution/v3/configuration"
 	"github.com/distribution/distribution/v3/registry"
 	_ "github.com/distribution/distribution/v3/registry/auth/htpasswd"
 	_ "github.com/distribution/distribution/v3/registry/storage/driver/inmemory"
-	"github.com/foxcpp/go-mockdns"
 	"github.com/stretchr/testify/suite"
 	"golang.org/x/crypto/bcrypt"
 
@@ -65,12 +63,10 @@ type TestSuite struct {
 	CompromisedRegistryHost string
 	WorkspaceDir            string
 	RegistryClient          *Client
-
-	// A mock DNS server needed for TLS connection testing.
-	srv *mockdns.Server
+	dockerRegistry          *registry.Registry
 }
 
-func setup(suite *TestSuite, tlsEnabled, insecure bool) *registry.Registry {
+func setup(suite *TestSuite, tlsEnabled, insecure bool) {
 	suite.WorkspaceDir = testWorkspaceDir
 	os.RemoveAll(suite.WorkspaceDir)
 	os.Mkdir(suite.WorkspaceDir, 0700)
@@ -135,13 +131,6 @@ func setup(suite *TestSuite, tlsEnabled, insecure bool) *registry.Registry {
 	// host is localhost/127.0.0.1.
 	port := ln.Addr().(*net.TCPAddr).Port
 	suite.DockerRegistryHost = fmt.Sprintf("helm-test-registry:%d", port)
-	suite.srv, err = mockdns.NewServer(map[string]mockdns.Zone{
-		"helm-test-registry.": {
-			A: []string{"127.0.0.1"},
-		},
-	}, false)
-	suite.Nil(err, "no error creating mock DNS server")
-	suite.srv.PatchNet(net.DefaultResolver)
 
 	config.HTTP.Addr = ln.Addr().String()
 	config.HTTP.DrainTimeout = time.Duration(10) * time.Second
@@ -166,20 +155,18 @@ func setup(suite *TestSuite, tlsEnabled, insecure bool) *registry.Registry {
 			config.HTTP.TLS.ClientCAs = []string{tlsCA}
 		}
 	}
-	dockerRegistry, err := registry.NewRegistry(context.Background(), config)
+	suite.dockerRegistry, err = registry.NewRegistry(context.Background(), config)
 	suite.Nil(err, "no error creating test registry")
 
 	suite.CompromisedRegistryHost = initCompromisedRegistryTestServer()
-	return dockerRegistry
+	go func() {
+		_ = suite.dockerRegistry.ListenAndServe()
+	}()
 }
 
 func teardown(suite *TestSuite) {
-	var lock sync.Mutex
-	lock.Lock()
-	defer lock.Unlock()
-	if suite.srv != nil {
-		mockdns.UnpatchNet(net.DefaultResolver)
-		suite.srv.Close()
+	if suite.dockerRegistry != nil {
+		_ = suite.dockerRegistry.Shutdown(context.Background())
 	}
 }