Use a key-based mutex lock to download resources/agent bins once.

juju · jujubot · Aug 6, 2021 · Aug 5, 2021 · Aug 6, 2021 · e309e883d79e8d9df196c4ec921ae5e95ae83c99
commit e309e883d79e8d9df196c4ec921ae5e95ae83c99
diff --git a/apiserver/tools.go b/apiserver/tools.go
@@ -16,10 +16,10 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/im7mortal/kmutex"
 	"github.com/juju/errors"
 	jujuhttp "github.com/juju/http/v2"
 	"github.com/juju/version/v2"
-	"golang.org/x/sync/singleflight"
 
 	"github.com/juju/juju/apiserver/common"
 	"github.com/juju/juju/apiserver/httpcontext"
@@ -67,13 +67,14 @@ type toolsUploadHandler struct {
 
 // toolsHandler handles tool download through HTTPS in the API server.
 type toolsDownloadHandler struct {
-	ctxt      httpContext
-	fetchOnce singleflight.Group
+	ctxt       httpContext
+	fetchMutex *kmutex.Kmutex
 }
 
 func newToolsDownloadHandler(httpCtxt httpContext) *toolsDownloadHandler {
 	return &toolsDownloadHandler{
-		ctxt: httpCtxt,
+		ctxt:       httpCtxt,
+		fetchMutex: kmutex.New(),
 	}
 }
 
@@ -215,6 +216,10 @@ func (h *toolsDownloadHandler) getToolsForRequest(r *http.Request, st *state.Sta
 		}
 	}
 
+	locker := h.fetchMutex.Locker(storageVers.String())
+	locker.Lock()
+	defer locker.Unlock()
+
 	md, reader, err := storage.Open(storageVers.String())
 	if errors.IsNotFound(err) {
 		// Tools could not be found in tools storage,
@@ -224,11 +229,8 @@ func (h *toolsDownloadHandler) getToolsForRequest(r *http.Request, st *state.Sta
 		if osTypeName != "" {
 			storageVers.Release = osTypeName
 		}
-		_, err, _ = h.fetchOnce.Do(storageVers.String(), func() (interface{}, error) {
-			return nil, h.fetchAndCacheTools(vers, storageVers)
-		})
+		err = h.fetchAndCacheTools(vers, storageVers)
 		if err != nil {
-			h.fetchOnce.Forget(storageVers.String())
 			err = errors.Annotate(err, "error fetching agent binaries")
 		} else {
 			md, reader, err = storage.Open(storageVers.String())

diff --git a/go.mod b/go.mod
@@ -44,6 +44,7 @@ require (
 	github.com/hashicorp/golang-lru v0.5.4 // indirect
 	github.com/hashicorp/raft v2.0.0-20200420012049-88ad3b3f0a54+incompatible
 	github.com/hashicorp/raft-boltdb v0.0.0-20171010151810-6e5ba93211ea
+	github.com/im7mortal/kmutex v1.0.1
 	github.com/imdario/mergo v0.3.10 // indirect
 	github.com/juju/ansiterm v0.0.0-20180109212912-720a0952cc2a
 	github.com/juju/blobstore/v2 v2.0.0-20210302045357-edd2b24570b7

diff --git a/go.sum b/go.sum
@@ -361,6 +361,8 @@ github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uG
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
+github.com/im7mortal/kmutex v1.0.1 h1:zAACzjwD+OEknDqnLdvRa/BhzFM872EBwKijviGLc9Q=
+github.com/im7mortal/kmutex v1.0.1/go.mod h1:f71c/Ugk/+58OHRAgvgzPP3QEiWGUjK13fd8ozfKWdo=
 github.com/imdario/mergo v0.3.5/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
 github.com/imdario/mergo v0.3.10 h1:6q5mVkdH/vYmqngx7kZQTjJ5HRsx+ImorDIEQ+beJgc=
 github.com/imdario/mergo v0.3.10/go.mod h1:jmQim1M+e3UYxmgPu/WyfjB3N3VflVyUjjjwH0dnCYA=

diff --git a/resource/repositories/mocks/mocks.go b/resource/repositories/mocks/mocks.go
diff --git a/resource/repositories/operations.go b/resource/repositories/operations.go
@@ -99,6 +99,10 @@ func GetResource(args GetResourceArgs) (resource.Resource, io.ReadCloser, error)
 		repo: args.Repository,
 	}
 
+	locker := opRepo.resourceLocker(args.Name)
+	locker.Lock()
+	defer locker.Unlock()
+
 	res, reader, err := opRepo.get(args.Name)
 	if err != nil {
 		return resource.Resource{}, nil, errors.Trace(err)
@@ -132,7 +136,6 @@ func GetResource(args GetResourceArgs) (resource.Resource, io.ReadCloser, error)
 	if err != nil {
 		return resource.Resource{}, nil, errors.Trace(err)
 	}
-
 	res, reader, err = opRepo.set(data.Resource, data, state.DoNotIncrementCharmModifiedVersion)
 	if err != nil {
 		return resource.Resource{}, nil, errors.Trace(err)

diff --git a/resource/repositories/operations_test.go b/resource/repositories/operations_test.go
@@ -0,0 +1,137 @@
+// Copyright 2021 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package repositories_test
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"sync"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/juju/charm/v8"
+	charmresource "github.com/juju/charm/v8/resource"
+	jc "github.com/juju/testing/checkers"
+	gc "gopkg.in/check.v1"
+
+	"github.com/juju/errors"
+	"github.com/juju/juju/charmstore"
+	"github.com/juju/juju/resource"
+	"github.com/juju/juju/resource/repositories"
+	"github.com/juju/juju/resource/repositories/mocks"
+	"github.com/juju/juju/state"
+)
+
+type OperationsSuite struct{}
+
+var _ = gc.Suite(&OperationsSuite{})
+
+func (s *OperationsSuite) TestConcurrentGetResource(c *gc.C) {
+	ctrl := gomock.NewController(c)
+	defer ctrl.Finish()
+	er := mocks.NewMockEntityRepository(ctrl)
+	rg := mocks.NewMockResourceGetter(ctrl)
+
+	stateLock := sync.Mutex{}
+	fetchMut := &sync.Mutex{}
+	fetchMut.Lock()
+
+	er.EXPECT().FetchLock(gomock.Any()).AnyTimes().Return(fetchMut)
+
+	openState := struct {
+		res    resource.Resource
+		buffer []byte
+		err    error
+	}{
+		err: errors.NotFoundf("resource"),
+	}
+	er.EXPECT().OpenResource(gomock.Any()).AnyTimes().DoAndReturn(func(name string) (resource.Resource, io.ReadCloser, error) {
+		stateLock.Lock()
+		defer stateLock.Unlock()
+		reader := io.ReadCloser(nil)
+		if openState.err == nil && openState.buffer != nil {
+			reader = ioutil.NopCloser(bytes.NewBuffer(openState.buffer))
+		}
+		return openState.res, reader, openState.err
+	})
+
+	getState := struct {
+		res resource.Resource
+	}{
+		res: resource.Resource{
+			ApplicationID: "gitlab",
+			Username:      "gitlab-0",
+			Resource: charmresource.Resource{
+				Meta: charmresource.Meta{
+					Name: "company-icon",
+				},
+				Origin: charmresource.OriginStore,
+			},
+		},
+	}
+	er.EXPECT().GetResource(gomock.Any()).AnyTimes().DoAndReturn(func(name string) (resource.Resource, error) {
+		stateLock.Lock()
+		defer stateLock.Unlock()
+		return getState.res, nil
+	})
+
+	gomock.InOrder(
+		rg.EXPECT().GetResource(repositories.ResourceRequest{
+			CharmID: repositories.CharmID{URL: charm.MustParseURL("cs:gitlab")},
+			Name:    "company-icon",
+		}).Times(1).Return(charmstore.ResourceData{
+			ReadCloser: ioutil.NopCloser(bytes.NewBufferString("data")),
+			Resource: charmresource.Resource{
+				Meta: charmresource.Meta{
+					Name: "company-icon",
+				},
+				Origin: charmresource.OriginStore,
+			},
+		}, nil),
+		er.EXPECT().SetResource(gomock.Any(), gomock.Any(), gomock.Any()).Times(1).DoAndReturn(func(res charmresource.Resource, reader io.Reader, arg state.IncrementCharmModifiedVersionType) (charmresource.Resource, error) {
+			stateLock.Lock()
+			defer stateLock.Unlock()
+			// Make sure this takes a while.
+			time.Sleep(10 * time.Millisecond)
+			buf, err := ioutil.ReadAll(reader)
+			if err != nil {
+				return charmresource.Resource{}, errors.Trace(err)
+			}
+			res.Size = int64(len(buf))
+			openState.buffer = buf
+			openState.err = nil
+			getState.res.Resource = res
+			openState.res = getState.res
+			return res, nil
+		}),
+	)
+
+	args := repositories.GetResourceArgs{
+		Client:     rg,
+		Repository: er,
+		Name:       "company-icon",
+		CharmID:    repositories.CharmID{URL: charm.MustParseURL("cs:gitlab")},
+	}
+
+	start := sync.WaitGroup{}
+	done := sync.WaitGroup{}
+	for i := 0; i < 100; i++ {
+		start.Add(1)
+		done.Add(1)
+		go func() {
+			defer done.Done()
+			start.Done()
+			rsc, reader, err := repositories.GetResource(args)
+			c.Check(err, jc.ErrorIsNil)
+			c.Check(reader, gc.NotNil)
+			c.Check(rsc, gc.DeepEquals, getState.res)
+		}()
+	}
+
+	start.Wait()
+	// This synchronises all the threads to start at the same time.
+	fetchMut.Unlock()
+	done.Wait()
+}
diff --git a/resource/repositories/package_test.go b/resource/repositories/package_test.go
@@ -0,0 +1,16 @@
+// Copyright 2021 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package repositories_test
+
+import (
+	"testing"
+
+	gc "gopkg.in/check.v1"
+)
+
+//go:generate go run github.com/golang/mock/mockgen -package mocks -destination mocks/mocks.go github.com/juju/juju/resource/repositories EntityRepository,ResourceGetter
+
+func Test(t *testing.T) {
+	gc.TestingT(t)
+}