Commit 07da5edb authored by Suleimi Ahmed's avatar Suleimi Ahmed 🔴 Committed by Hayley Swimelar
Browse files

fix(handlers): prevent fk violation cache race

parent 51de46d6

registry/datastore/repository.go

+45 −11
Original line number Diff line number Diff line
@@ -136,6 +136,7 @@ type repositoryStore struct { 
	// db can be either a *sql.DB or *sql.Tx

	db                  Queryer

	cache               RepositoryCache

	disableTxCacheWrite bool

}



// NewRepositoryStore builds a new repositoryStore.
@@ -145,10 +146,43 @@ func NewRepositoryStore(db Queryer, opts ...RepositoryStoreOption) RepositorySto 
	for _, o := range opts {

		o(rStore)

	}

	// If the database is a transaction, disable cache writes through the repositoryStore.

	// This is to avoid data inconsistencies and race conditions when other

	// concurrent queries or requests retrieve cached objects/rows before the

	// creating transaction commits to the database which should be the single source of truth.

	if _, isTransaction := db.(*Tx); isTransaction {

		rStore.disableTxCacheWrite = true

	}



	return rStore

}



// cacheSet sets the repository in the cache if cache writes are not disabled.

func (s *repositoryStore) cacheSet(ctx context.Context, r *models.Repository) {

	switch {

	case s.disableTxCacheWrite:

		log.GetLogger(log.WithContext(ctx)).Debug("cacheSet: cache writes are not allowed in a database transaction")

		return

	case s.cache == nil:

		return

	default:

		s.cache.Set(ctx, r)

	}

}



// cacheInvalidateSize invalidates the repository size in the cache if cache writes are not disabled.

func (s *repositoryStore) cacheInvalidateSize(ctx context.Context, r *models.Repository) {

	switch {

	case s.disableTxCacheWrite:

		log.GetLogger(log.WithContext(ctx)).Debug("cacheInvalidateSize: cache writes are not allowed in a database transaction")

		return

	case s.cache == nil:

		return

	default:

		s.cache.InvalidateSize(ctx, r)

	}

}



// RepositoryManifestService implements the validation.ManifestExister

// interface for repository-scoped manifests.

type RepositoryManifestService struct {
@@ -616,7 +650,7 @@ func (s *repositoryStore) FindByPath(ctx context.Context, path string) (*models. 
		return r, err

	}



	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return r, nil

}
@@ -1721,7 +1755,7 @@ func (s *repositoryStore) Create(ctx context.Context, r *models.Repository) erro 
		return fmt.Errorf("creating repository: %w", err)

	}



	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return nil

}
@@ -1800,7 +1834,7 @@ func (s *repositoryStore) Size(ctx context.Context, r *models.Repository) (Repos 


	// Update the size attribute for the cached repository object

	r.Size = &b

	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return RepositorySize{bytes: b}, nil

}
@@ -2083,7 +2117,7 @@ func (s *repositoryStore) CreateOrFind(ctx context.Context, r *models.Repository 
			return err

		}

		*r = *tmp

		s.cache.Set(ctx, r)

		s.cacheSet(ctx, r)

	}



	return nil
@@ -2122,7 +2156,7 @@ func (s *repositoryStore) CreateByPath(ctx context.Context, path string, opts .. 
		return nil, err

	}



	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return r, nil

}
@@ -2151,7 +2185,7 @@ func (s *repositoryStore) CreateOrFindByPath(ctx context.Context, path string, o 
		return nil, err

	}



	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return r, nil

}
@@ -2177,7 +2211,7 @@ func (s *repositoryStore) Update(ctx context.Context, r *models.Repository) erro 
		return fmt.Errorf("updating repository: %w", err)

	}



	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return nil

}
@@ -2240,7 +2274,7 @@ func (s *repositoryStore) DeleteTagByName(ctx context.Context, r *models.Reposit 
		return false, fmt.Errorf("deleting tag: %w", err)

	}



	s.cache.InvalidateSize(ctx, r)

	s.cacheInvalidateSize(ctx, r)



	return count == 1, nil

}
@@ -2271,7 +2305,7 @@ func (s *repositoryStore) DeleteManifest(ctx context.Context, r *models.Reposito 
		return false, fmt.Errorf("deleting manifest: %w", err)

	}



	s.cache.InvalidateSize(ctx, r)

	s.cacheInvalidateSize(ctx, r)



	return count == 1, nil

}
@@ -2345,7 +2379,7 @@ func (s *repositoryStore) Rename(ctx context.Context, r *models.Repository, newP 
		return fmt.Errorf("renaming repository: %w", err)

	}



	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return nil

}
@@ -2377,7 +2411,7 @@ func (s *repositoryStore) UpdateLastPublishedAt(ctx context.Context, r *models.R 
		return fmt.Errorf("updating repository last published at: %w", err)

	}



	s.cache.Set(ctx, r)

	s.cacheSet(ctx, r)



	return nil

}

registry/datastore/repository_integration_test.go

+0 −34
Original line number Diff line number Diff line
@@ -1962,40 +1962,6 @@ func TestRepositoryStore_SizeWithDescendants_TopLevel_ChecksCacheForPreviousTime 
	require.Zero(t, size.Bytes())

}



func TestRepositoryStore_SizeWithDescendants_TopLevel_SetsCacheOnTimeout(t *testing.T) {

	reloadManifestFixtures(t)



	redisCache, redisMock := itestutil.RedisCacheMock(t, 0)

	cache := datastore.NewCentralRepositoryCache(redisCache)



	// use transaction with a statement timeout of 1ms, so that all queries within time out

	tx, err := suite.db.BeginTx(suite.ctx, nil)

	require.NoError(t, err)

	defer tx.Rollback()



	_, err = tx.ExecContext(suite.ctx, "SET statement_timeout TO 1")

	require.NoError(t, err)

	// wait a bit so that PG has time to flush the update

	time.Sleep(250 * time.Millisecond)



	s := datastore.NewRepositoryStore(tx, datastore.WithRepositoryCache(cache))



	repo := &models.Repository{NamespaceID: 3, ID: 8, Path: "usage-group"}

	redisKey := fmt.Sprintf("registry:db:{repository:%s:%s}:swd-timeout", repo.Path, digest.FromString(repo.Path).Hex())



	redisMock.ExpectGet(redisKey).RedisNil()

	redisMock.ExpectSet(redisKey, "true", 24*time.Hour).SetVal("true")



	size, err := s.SizeWithDescendants(suite.ctx, repo)

	require.Error(t, err)



	// make sure the error is not masked

	var pgErr *pgconn.PgError

	require.ErrorAs(t, err, &pgErr)

	require.Equal(t, pgerrcode.QueryCanceled, pgErr.Code)

	require.Zero(t, size)

}



func TestRepositoryStore_SizeWithDescendants_NonTopLevel_DoesNotTouchCacheTimeout(t *testing.T) {

	reloadManifestFixtures(t)

registry/datastore/repository_internal_test.go

+71 −0
Original line number Diff line number Diff line 
package datastore



import (

	"context"

	"database/sql"

	"fmt"

	"strings"

	"testing"

	"time"



	"github.com/DATA-DOG/go-sqlmock"

	"github.com/docker/distribution/registry/datastore/models"

	itestutil "github.com/docker/distribution/registry/internal/testutil"

	"github.com/jackc/pgerrcode"

	"github.com/jackc/pgx/v5/pgconn"

	"github.com/opencontainers/go-digest"

	"github.com/stretchr/testify/require"

)
@@ -242,3 +250,66 @@ func Test_tagsDetailPaginatedQuery(t *testing.T) { 
		})

	}

}



type dbMock struct {

	db *sql.DB

}



func (m *dbMock) QueryRowContext(ctx context.Context, query string, args ...any) *Row {

	sqlRow := m.db.QueryRowContext(ctx, query, args...)

	return &Row{

		Row: sqlRow,

	}

}



func (m *dbMock) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {

	return m.db.ExecContext(ctx, query, args...)

}



func (m *dbMock) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {

	return m.db.QueryContext(ctx, query, args...)

}



func TestRepositoryStore_SizeWithDescendants_TopLevel_SetsCacheOnTimeout(t *testing.T) {

	db, mock, err := sqlmock.New()

	require.NoError(t, err)

	defer db.Close()



	redisCache, redisMock := itestutil.RedisCacheMock(t, 0)

	cache := NewCentralRepositoryCache(redisCache)



	// The error we want Postgres to return is a statement timeout cancellation error

	pgTimeout := &pgconn.PgError{

		Code:    pgerrcode.QueryCanceled,

		Message: "canceling statement due to statement timeout",

	}



	// Set up redis expectations

	repo := &models.Repository{NamespaceID: 3, ID: 8, Path: "usage-group"}



	// Mock the SQL query SizeWithDescendants()

	mock.ExpectQuery(`(?s)SELECT\s+coalesce\(sum\(q\.size\), 0\).*WITH RECURSIVE cte`).

		WithArgs(repo.NamespaceID).

		WillReturnError(pgTimeout)



	redisKey := fmt.Sprintf(

		"registry:db:{repository:%s:%s}:swd-timeout",

		repo.Path,

		digest.FromString(repo.Path).Hex(),

	)



	redisMock.ExpectGet(redisKey).RedisNil()

	redisMock.ExpectSet(redisKey, "true", 24*time.Hour).SetVal("true")



	s := NewRepositoryStore(&dbMock{db: db}, WithRepositoryCache(cache))

	size, err := s.SizeWithDescendants(context.Background(), repo)

	require.Error(t, err)



	var pgErr *pgconn.PgError

	require.ErrorAs(t, err, &pgErr)

	require.Equal(t, pgerrcode.QueryCanceled, pgErr.Code)

	require.Zero(t, size)



	// Ensure all db expectations ran

	require.NoError(t, mock.ExpectationsWereMet())

}

registry/handlers/blobupload.go

+5 −0
Original line number Diff line number Diff line
@@ -248,6 +248,11 @@ func dbPutBlobUploadComplete(ctx context.Context, db *datastore.DB, repoPath str 
		return fmt.Errorf("committing database transaction: %w", err)

	}



	// Set the repository in the cache outside of the transaction.

	if repoCache != nil {

		repoCache.Set(ctx, r)

	}



	return nil

}

registry/handlers/blobupload_integration_test.go

+0 −5
Original line number Diff line number Diff line
@@ -257,7 +257,6 @@ func TestDBPutBlobUploadComplete_NonExistantRepoAndBlob(t *testing.T) { 
	repoCacheMock := mocks.NewMockRepositoryCache(ctrl)

	gomock.InOrder(

		repoCacheMock.EXPECT().Get(env.ctx, repoName).Return(nil).Times(3),

		repoCacheMock.EXPECT().Set(env.ctx, (*models.Repository)(nil)).Times(1),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

	)
@@ -287,7 +286,6 @@ func TestDBPutBlobUploadComplete_BlobExistsAndNonExistentRepo(t *testing.T) { 
	repoCacheMock := mocks.NewMockRepositoryCache(ctrl)

	gomock.InOrder(

		repoCacheMock.EXPECT().Get(env.ctx, repoName).Return(nil).Times(3),

		repoCacheMock.EXPECT().Set(env.ctx, (*models.Repository)(nil)).Times(1),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

	)
@@ -316,7 +314,6 @@ func TestDBPutBlobUploadComplete_RepoExistsAndBlobDoesNot(t *testing.T) { 
	gomock.InOrder(

		repoCacheMock.EXPECT().Get(env.ctx, repoName).Return(nil).Times(3),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

	)



	err := dbPutBlobUploadComplete(env.ctx, env.db, repoName, desc, repoCacheMock)
@@ -346,7 +343,6 @@ func TestDBPutBlobUploadComplete_BothBlobAndRepoExistsButNotLinked(t *testing.T) 
	gomock.InOrder(

		repoCacheMock.EXPECT().Get(env.ctx, repoName).Return(nil).Times(3),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

	)



	err := dbPutBlobUploadComplete(env.ctx, env.db, repoName, desc, repoCacheMock)
@@ -375,7 +371,6 @@ func TestDBPutBlobUploadComplete_BothBlobAndRepoExistsAndLinked(t *testing.T) { 
	gomock.InOrder(

		repoCacheMock.EXPECT().Get(env.ctx, repoName).Return(nil).Times(3),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

		repoCacheMock.EXPECT().Set(env.ctx, gomock.Cond(gomockMatchRepoFn(repoName))).Times(1),

	)



	err := dbPutBlobUploadComplete(env.ctx, env.db, repoName, desc, repoCacheMock)