fix(build-index): Add retry logic for manifest download during dependency resolution (#502)

* fix(build-index): Add retry logic for manifest download during dependency resolution

* Fix

* Fix

* Fix

* Revert log

* Test retries

* Increase timeout, fix

Author: hweawer

build-index/tagserver/server.go

@@ -180,15 +180,15 @@ func (s *Server) putTagHandler(w http.ResponseWriter, r *http.Request) error {
 	}
 	replicate, err := strconv.ParseBool(httputil.GetQueryArg(r, "replicate", "false"))
 	if err != nil {
-		return handler.Errorf("parse query arg `replicate`: %s", err)
+		return fmt.Errorf("parse query arg `replicate`: %w", err)
 	}
 
 	log.With("tag", tag, "digest", d.String(), "replicate", replicate).Info("Putting tag")
 
 	deps, err := s.depResolver.Resolve(tag, d)
 	if err != nil {
 		log.With("tag", tag, "digest", d.String(), "error", err).Error("Failed to resolve dependencies")
-		return fmt.Errorf("resolve dependencies: %s", err)
+		return fmt.Errorf("resolve dependencies: %w", err)
 	}
 
 	log.With("tag", tag, "digest", d.String(), "dependency_count", len(deps)).Debug("Resolved dependencies")
@@ -252,7 +252,6 @@ func (s *Server) getTagHandler(w http.ResponseWriter, r *http.Request) error {
 	d, err := s.store.Get(tag)
 	if err != nil {
 		if err == tagstore.ErrTagNotFound {
-			log.With("tag", tag).Debug("Tag not found")
 			return handler.ErrorStatus(http.StatusNotFound)
 		}
 		log.With("tag", tag).Errorf("Failed to get tag from storage: %s", err)
@@ -465,21 +464,18 @@ func (s *Server) getOriginHandler(w http.ResponseWriter, r *http.Request) error
 func (s *Server) putTag(tag string, d core.Digest, deps core.DigestList) error {
 	log.With("tag", tag, "digest", d.String(), "dependency_count", len(deps)).Debug("Validating tag dependencies")
 
-	for i, dep := range deps {
+	for _, dep := range deps {
 		if _, err := s.localOriginClient.Stat(tag, dep); err == blobclient.ErrBlobNotFound {
-			log.With("tag", tag, "digest", d.String(), "missing_dependency", dep.String(), "dependency_index", i).Error("Missing dependency blob")
-			return handler.Errorf("cannot upload tag, missing dependency %s", dep)
+			return fmt.Errorf("cannot upload tag, missing dependency %s", dep)
 		} else if err != nil {
-			log.With("tag", tag, "digest", d.String(), "dependency", dep.String(), "dependency_index", i).Errorf("Failed to check dependency blob: %s", err)
-			return handler.Errorf("check blob: %s", err)
+			return fmt.Errorf("check blob: %w", err)
 		}
 	}
 
 	log.With("tag", tag, "digest", d.String()).Debug("All dependencies validated successfully")
 
 	if err := s.store.Put(tag, d, 0); err != nil {
-		log.With("tag", tag, "digest", d.String(), "error", err).Error("Failed to store tag")
-		return handler.Errorf("storage: %s", err)
+		return fmt.Errorf("storage: %w", err)
 	}
 
 	log.With("tag", tag, "digest", d.String()).Info("Tag stored locally")
@@ -526,8 +522,7 @@ func (s *Server) replicateTag(tag string, d core.Digest, deps core.DigestList) e
 	for _, dest := range destinations {
 		task := tagreplication.NewTask(tag, d, deps, dest, 0)
 		if err := s.tagReplicationManager.Add(task); err != nil {
-			log.With("tag", tag, "digest", d.String(), "destination", dest).Errorf("Failed to add remote replication task: %s", err)
-			return handler.Errorf("add replicate task: %s", err)
+			return fmt.Errorf("add replicate task: %w", err)
 		}
 		log.With("tag", tag, "digest", d.String(), "destination", dest).Debug("Added remote replication task")
 	}

build-index/tagtype/docker_resolver.go

@@ -17,37 +17,66 @@ import (
 	"bytes"
 	"fmt"
 
+	"github.com/cenkalti/backoff"
 	"github.com/docker/distribution"
+
 	"github.com/uber/kraken/core"
 	"github.com/uber/kraken/origin/blobclient"
 	"github.com/uber/kraken/utils/dockerutil"
+	"github.com/uber/kraken/utils/httputil"
+	"github.com/uber/kraken/utils/log"
 )
 
 type dockerResolver struct {
-	originClient blobclient.ClusterClient
+	originClient  blobclient.ClusterClient
+	backoffConfig httputil.ExponentialBackOffConfig
 }
 
 // Resolve returns all layers + manifest of given tag as its dependencies.
 func (r *dockerResolver) Resolve(tag string, d core.Digest) (core.DigestList, error) {
 	m, err := r.downloadManifest(tag, d)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("download manifest: %w", err)
 	}
 	deps, err := dockerutil.GetManifestReferences(m)
 	if err != nil {
-		return nil, fmt.Errorf("get manifest references: %s", err)
+		return nil, fmt.Errorf("get manifest references: %w", err)
 	}
 	return append(deps, d), nil
 }
 
 func (r *dockerResolver) downloadManifest(tag string, d core.Digest) (distribution.Manifest, error) {
 	buf := &bytes.Buffer{}
-	if err := r.originClient.DownloadBlob(tag, d, buf); err != nil {
-		return nil, fmt.Errorf("download blob: %s", err)
+	attempt := 0
+
+	retryFunc := func() error {
+		attempt++
+		buf.Reset()
+
+		err := r.originClient.DownloadBlob(tag, d, buf)
+		if err == nil {
+			return nil
+		}
+
+		if attempt > 1 {
+			log.With("tag", tag, "digest", d.Hex(), "attempt", attempt, "error", err).
+				Warn("Manifest download failed, will retry")
+		}
+
+		if err != blobclient.ErrBlobNotFound &&
+			!httputil.IsNetworkError(err) {
+			return backoff.Permanent(err)
+		}
+
+		return err
+	}
+
+	if err := backoff.Retry(retryFunc, r.backoffConfig.Build()); err != nil {
+		return nil, err
 	}
 	manifest, _, err := dockerutil.ParseManifest(buf)
 	if err != nil {
-		return nil, fmt.Errorf("parse manifest: %s", err)
+		return nil, fmt.Errorf("parse manifest: %w", err)
 	}
 	return manifest, nil
 }

build-index/tagtype/docker_resolver_test.go

@@ -0,0 +1,240 @@
+// Copyright (c) 2016-2019 Uber Technologies, Inc.
+//
+// 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 tagtype
+
+import (
+	"io"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/uber/kraken/core"
+	mockblobclient "github.com/uber/kraken/mocks/origin/blobclient"
+	"github.com/uber/kraken/origin/blobclient"
+	"github.com/uber/kraken/utils/dockerutil"
+	"github.com/uber/kraken/utils/httputil"
+	"github.com/uber/kraken/utils/mockutil"
+)
+
+func newTestDockerResolver(ctrl *gomock.Controller) (*dockerResolver, *mockblobclient.MockClusterClient) {
+	originClient := mockblobclient.NewMockClusterClient(ctrl)
+
+	// Use minimal backoff for fast tests
+	backoffConfig := httputil.ExponentialBackOffConfig{
+		Enabled:             true,
+		InitialInterval:     1 * time.Millisecond,
+		RandomizationFactor: 0,
+		Multiplier:          2,
+		MaxInterval:         10 * time.Millisecond,
+		MaxRetries:          3,
+	}
+
+	resolver := &dockerResolver{
+		originClient:  originClient,
+		backoffConfig: backoffConfig,
+	}
+
+	return resolver, originClient
+}
+
+// setupDockerResolverTest sets up common test dependencies.
+// Returns require, ctrl, resolver, and mockOrigin.
+func setupDockerResolverTest(t *testing.T) (*require.Assertions, *gomock.Controller, *dockerResolver, *mockblobclient.MockClusterClient) {
+	require := require.New(t)
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
+
+	resolver, mockOrigin := newTestDockerResolver(ctrl)
+
+	return require, ctrl, resolver, mockOrigin
+}
+
+// setupDockerResolverTestWithManifest sets up common test dependencies with manifest fixtures.
+// Returns require, ctrl, resolver, mockOrigin, tag, layers, manifest digest, and manifest bytes.
+func setupDockerResolverTestWithManifest(t *testing.T) (*require.Assertions, *gomock.Controller, *dockerResolver, *mockblobclient.MockClusterClient, string, core.DigestList, core.Digest, []byte) {
+	require, ctrl, resolver, mockOrigin := setupDockerResolverTest(t)
+
+	tag := "repo/image:v1.0"
+	layers := core.DigestListFixture(3)
+	manifest, manifestBytes := dockerutil.ManifestFixture(layers[0], layers[1], layers[2])
+
+	return require, ctrl, resolver, mockOrigin, tag, layers, manifest, manifestBytes
+}
+
+func TestDockerResolver_DownloadManifest_Success(t *testing.T) {
+	require, _, resolver, mockOrigin, tag, _, manifest, manifestBytes := setupDockerResolverTestWithManifest(t)
+
+	// Expect successful download on first attempt
+	mockOrigin.EXPECT().
+		DownloadBlob(tag, manifest, mockutil.MatchWriter(manifestBytes)).
+		Return(nil)
+
+	result, err := resolver.downloadManifest(tag, manifest)
+	require.NoError(err)
+	require.NotNil(result)
+}
+
+func TestDockerResolver_DownloadManifest_RetryOnBlobNotFound(t *testing.T) {
+	require, _, resolver, mockOrigin, tag, _, manifest, manifestBytes := setupDockerResolverTestWithManifest(t)
+
+	// First attempt fails with blob not found, second succeeds
+	gomock.InOrder(
+		mockOrigin.EXPECT().
+			DownloadBlob(tag, manifest, gomock.Any()).
+			Return(blobclient.ErrBlobNotFound),
+		mockOrigin.EXPECT().
+			DownloadBlob(tag, manifest, mockutil.MatchWriter(manifestBytes)).
+			Return(nil),
+	)
+
+	result, err := resolver.downloadManifest(tag, manifest)
+	require.NoError(err)
+	require.NotNil(result)
+}
+
+func TestDockerResolver_DownloadManifest_ExhaustedRetries(t *testing.T) {
+	require, _, resolver, mockOrigin := setupDockerResolverTest(t)
+
+	tag := "repo/image:v1.0"
+	manifest := core.DigestFixture()
+
+	// All attempts fail with blob not found (MaxRetries=3, so 4 total attempts)
+	mockOrigin.EXPECT().
+		DownloadBlob(tag, manifest, gomock.Any()).
+		Return(blobclient.ErrBlobNotFound).
+		Times(4)
+
+	result, err := resolver.downloadManifest(tag, manifest)
+	require.Error(err)
+	require.Nil(result)
+	require.Equal(blobclient.ErrBlobNotFound, err)
+}
+
+func TestDockerResolver_DownloadManifest_PermanentError(t *testing.T) {
+	require, _, resolver, mockOrigin := setupDockerResolverTest(t)
+
+	tag := "repo/image:v1.0"
+	manifest := core.DigestFixture()
+
+	// 401 Unauthorized is a permanent error (4xx except 404)
+	permanentErr := httputil.StatusError{Status: 401}
+
+	mockOrigin.EXPECT().
+		DownloadBlob(tag, manifest, gomock.Any()).
+		Return(permanentErr).
+		Times(1) // Should only be called once, no retries
+
+	result, err := resolver.downloadManifest(tag, manifest)
+	require.Error(err)
+	require.Nil(result)
+}
+
+func TestDockerResolver_DownloadManifest_BufferResetBetweenRetries(t *testing.T) {
+	require, _, resolver, mockOrigin, tag, _, manifest, manifestBytes := setupDockerResolverTestWithManifest(t)
+
+	partialData := []byte("partial corrupt data")
+
+	// First attempt writes partial data then fails
+	// Second attempt succeeds with full data
+	gomock.InOrder(
+		mockOrigin.EXPECT().
+			DownloadBlob(tag, manifest, gomock.Any()).
+			DoAndReturn(func(tag string, d core.Digest, dst io.Writer) error {
+				dst.Write(partialData)
+				return blobclient.ErrBlobNotFound
+			}),
+		mockOrigin.EXPECT().
+			DownloadBlob(tag, manifest, mockutil.MatchWriter(manifestBytes)).
+			Return(nil),
+	)
+
+	result, err := resolver.downloadManifest(tag, manifest)
+	require.NoError(err)
+	require.NotNil(result)
+	// If buffer wasn't reset, parsing would fail due to partial data
+}
+
+func TestDockerResolver_DownloadManifest_InvalidManifestFormat(t *testing.T) {
+	require, _, resolver, mockOrigin := setupDockerResolverTest(t)
+
+	tag := "repo/image:v1.0"
+	manifest := core.DigestFixture()
+
+	// Download succeeds but returns invalid manifest data
+	mockOrigin.EXPECT().
+		DownloadBlob(tag, manifest, gomock.Any()).
+		DoAndReturn(func(tag string, d core.Digest, dst io.Writer) error {
+			dst.Write([]byte("invalid manifest json"))
+			return nil
+		})
+
+	result, err := resolver.downloadManifest(tag, manifest)
+	require.Error(err)
+	require.Nil(result)
+	require.Contains(err.Error(), "parse manifest")
+}
+
+func TestDockerResolver_Resolve_Success(t *testing.T) {
+	require, _, resolver, mockOrigin, tag, layers, manifest, manifestBytes := setupDockerResolverTestWithManifest(t)
+
+	mockOrigin.EXPECT().
+		DownloadBlob(tag, manifest, mockutil.MatchWriter(manifestBytes)).
+		Return(nil)
+
+	deps, err := resolver.Resolve(tag, manifest)
+	require.NoError(err)
+	require.Equal(core.DigestList(append(layers, manifest)), deps)
+}
+
+func TestDockerResolver_Resolve_DownloadError(t *testing.T) {
+	require, _, resolver, mockOrigin := setupDockerResolverTest(t)
+
+	tag := "repo/image:v1.0"
+	manifest := core.DigestFixture()
+
+	// All retries exhausted
+	mockOrigin.EXPECT().
+		DownloadBlob(tag, manifest, gomock.Any()).
+		Return(blobclient.ErrBlobNotFound).
+		Times(4)
+
+	deps, err := resolver.Resolve(tag, manifest)
+	require.Error(err)
+	require.Nil(deps)
+	require.Contains(err.Error(), "download manifest")
+}
+
+func TestDockerResolver_Resolve_WithRetries(t *testing.T) {
+	require, _, resolver, mockOrigin, tag, layers, manifest, manifestBytes := setupDockerResolverTestWithManifest(t)
+
+	// Fails twice, succeeds on third attempt
+	gomock.InOrder(
+		mockOrigin.EXPECT().
+			DownloadBlob(tag, manifest, gomock.Any()).
+			Return(blobclient.ErrBlobNotFound),
+		mockOrigin.EXPECT().
+			DownloadBlob(tag, manifest, gomock.Any()).
+			Return(blobclient.ErrBlobNotFound),
+		mockOrigin.EXPECT().
+			DownloadBlob(tag, manifest, mockutil.MatchWriter(manifestBytes)).
+			Return(nil),
+	)
+
+	deps, err := resolver.Resolve(tag, manifest)
+	require.NoError(err)
+	require.NotNil(deps)
+	require.Equal(core.DigestList(append(layers, manifest)), deps)
+}