internal/postgres/requeue.go - pkgsite - Git at Google

 // Copyright 2020 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package postgres

 import (
 	"context"
 	"database/sql"
 	"fmt"
 	"net/http"
 	"strconv"

 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/config"
 	"golang.org/x/pkgsite/internal/derrors"
 	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/log"
 )

 // UpdateModuleVersionStatesForReprocessing marks modules to be reprocessed
 // that were processed prior to the provided appVersion.
 func (db *DB) UpdateModuleVersionStatesForReprocessing(ctx context.Context, appVersion string) (err error) {
 	defer derrors.Wrap(&err, "UpdateModuleVersionStatesForReprocessing(ctx, %q)", appVersion)

 	for _, status := range []int{
 		http.StatusOK,
 		derrors.ToStatus(derrors.HasIncompletePackages),
 		derrors.ToStatus(derrors.DBModuleInsertInvalid),
 	} {
 		if err := db.UpdateModuleVersionStatesWithStatus(ctx, status, appVersion); err != nil {
 			return err
 		}
 	}
 	return nil
 }

 func (db *DB) UpdateModuleVersionStatesWithStatus(ctx context.Context, status int, appVersion string) (err error) {
 	query := `UPDATE module_version_states
 			SET
 				status = $2,
 				next_processed_after = CURRENT_TIMESTAMP,
 				last_processed_at = NULL
 			WHERE
 				app_version < $1
 				AND status = $3;`
 	affected, err := db.db.Exec(ctx, query, appVersion, derrors.ToReprocessStatus(status), status)
 	if err != nil {
 		return err
 	}
 	log.Infof(ctx,
 		"Updated module_version_states with status=%d and app_version < %q to status=%d; %d affected",
 		status, appVersion, derrors.ToReprocessStatus(status), affected)
 	return nil
 }

 // largeModulePackageThresold represents the package threshold at which it
 // becomes difficult to process packages. Modules with more than this number
 // of packages are generally different versions or forks of kubernetes,
 // aws-sdk-go, azure-sdk-go, and bilibili.
 const largeModulePackageThreshold = 1500

 // largeModulesLimit represents the number of large modules that we are
 // willing to enqueue at a given time.
 // var for testing.
 var largeModulesLimit = config.GetEnvInt("GO_DISCOVERY_LARGE_MODULES_LIMIT", 100)

 // GetNextModulesToFetch returns the next batch of modules that need to be
 // processed. We prioritize modules based on (1) whether it has status zero
 // (never processed), (2) whether it is the latest version, (3) if it is an
 // alternative module, and (4) the number of packages it has. We want to leave
 // time-consuming modules until the end and process them at a slower rate to
 // reduce database load and timeouts. We also want to leave alternative modules
 // towards the end, since these will incur unnecessary deletes otherwise.
 func (db *DB) GetNextModulesToFetch(ctx context.Context, limit int) (_ []*internal.ModuleVersionState, err error) {
 	defer derrors.Wrap(&err, "GetNextModulesToFetch(ctx, %d)", limit)
 	queryFmt := nextModulesToProcessQuery
 	if experiment.IsActive(ctx, internal.ExperimentAltRequeue) {
 		log.Infof(ctx, "using alternative requeue query")
 		queryFmt = nextModulesToProcessQueryAlt
 	}

 	var mvs []*internal.ModuleVersionState
 	query := fmt.Sprintf(queryFmt, moduleVersionStateColumns)

 	collect := func(rows *sql.Rows) error {
 		// Scan the last two columns separately; they are in the query only for sorting.
 		scan := func(dests ...interface{}) error {
 			var (
 				latest bool
 				npkg   int
 			)
 			return rows.Scan(append(dests, &latest, &npkg)...)
 		}
 		mv, err := scanModuleVersionState(scan)
 		if err != nil {
 			return err
 		}
 		mvs = append(mvs, mv)
 		return nil
 	}
 	if err := db.db.RunQuery(ctx, query, collect, largeModulePackageThreshold, limit); err != nil {
 		return nil, err
 	}
 	if len(mvs) == 0 {
 		log.Infof(ctx, "No modules to requeue")
 	} else {
 		fmtIntp := func(p *int) string {
 			if p == nil {
 				return "NULL"
 			}
 			return strconv.Itoa(*p)
 		}
 		start := mvs[0]
 		end := mvs[len(mvs)-1]
 		pkgRange := fmt.Sprintf("%s <= num_packages <= %s", fmtIntp(start.NumPackages), fmtIntp(end.NumPackages))
 		log.Infof(ctx, "GetNextModulesToFetch: num_modules=%d; %s; start_module=%q; end_module=%q",
 			len(mvs), pkgRange,
 			fmt.Sprintf("%s/@v/%s", start.ModulePath, start.Version),
 			fmt.Sprintf("%s/@v/%s", end.ModulePath, end.Version))
 	}

 	// Don't return more than largeModulesLimit of modules that have more than
 	// largeModulePackageThreshold packages, or of modules with status zero.
 	nLargeOrZero := 0
 	for i, m := range mvs {
 		if m.Status == 0 || (m.NumPackages != nil && *m.NumPackages >= largeModulePackageThreshold) {
 			nLargeOrZero++
 		}
 		if nLargeOrZero > largeModulesLimit {
 			return mvs[:i], nil
 		}
 	}
 	return mvs, nil
 }

 const nextModulesToProcessQuery = `
     -- Make a table of the latest versions of each module.
 	WITH latest_versions AS (
 		SELECT DISTINCT ON (module_path) module_path, version
 		FROM module_version_states
 		ORDER BY
 			module_path,
 			incompatible,
 			right(sort_version, 1) = '~' DESC, -- prefer release versions
 			sort_version DESC
 	)
 	SELECT %s, latest, npkg
 	FROM (
 		SELECT
 			%[1]s,
 			((module_path, version) IN (SELECT * FROM latest_versions)) AS latest,
 			COALESCE(num_packages, 0) AS npkg
 		FROM module_version_states
 	) s
 	WHERE next_processed_after < CURRENT_TIMESTAMP
 		AND (status = 0 OR status >= 500)
 	ORDER BY
 		CASE
 			 -- new modules
 			 WHEN status = 0 THEN 0
 			 WHEN npkg < $1 THEN		-- non-large modules
 				CASE
 					WHEN latest THEN	-- latest version
 						CASE
 							-- with ReprocessStatusOK or ReprocessHasIncompletePackages
 							WHEN status = 520 OR status = 521 THEN 1
 							-- with ReprocessBadModule or ReprocessAlternative or ReprocessDBModuleInsertInvalid
 							WHEN status = 540 OR status = 541 OR status = 542 THEN 2
 							ELSE 9
 						END
 					ELSE				-- non-latest version
 						CASE
 							WHEN status = 520 OR status = 521 THEN 3
 							WHEN status = 540 OR status = 541 OR status = 542 THEN 4
 							ELSE 9
 						END
 				END
 			ELSE						-- large modules
 				CASE
 					WHEN latest THEN	-- latest version
 						CASE
 							WHEN status = 520 OR status = 521 THEN 5
 							WHEN status = 540 OR status = 541 OR status = 542 THEN 6
 							ELSE 9
 						END
 					ELSE				-- non-latest version
 						CASE
 							WHEN status = 520 OR status = 521 THEN 7
 							WHEN status = 540 OR status = 541 OR status = 542 THEN 8
 							ELSE 9
 						END
 				END
 		END,
 		npkg,  -- prefer fewer packages
 		module_path,
 		version -- for reproducibility
 	LIMIT $2
 `

 // Like the standard query above, this query prioritizes latest versions.
 // However, it does not hold off large modules until the end. In fact, it tries
 // to avoid grouping modules in any way except by latest and status code:
 // processing is much smoother when they are enqueued in random order.
 //
 // To make the result deterministic for testing, we hash the module path and version
 // rather than actually choosing a random number. md5 is built in to postgres and
 // is an adequate hash for this purpose.
 const nextModulesToProcessQueryAlt = `
     -- Make a table of the latest versions of each module.
 	WITH latest_versions AS (
 		SELECT DISTINCT ON (module_path) module_path, version
 		FROM module_version_states
 		ORDER BY
 			module_path,
 			incompatible,
 			right(sort_version, 1) = '~' DESC, -- prefer release versions
 			sort_version DESC
 	)
 	SELECT %s, latest, npkg
 	FROM (
 		SELECT
 			%[1]s,
 			((module_path, version) IN (SELECT * FROM latest_versions)) AS latest,
 			COALESCE(num_packages, 0) AS npkg
 		FROM module_version_states
 	) s
 	WHERE next_processed_after < CURRENT_TIMESTAMP
 		AND (status = 0 OR status >= 500)
         AND $1 = $1 -- ignore largeModulePackageThreshold
 	ORDER BY
 		CASE
 			-- new modules
 			WHEN status = 0 THEN 0
 			WHEN latest THEN
 				CASE
 					-- with SheddingLoad or ReprocessStatusOK or ReprocessHasIncompletePackages
 					WHEN status = 503 or status = 520 OR status = 521 THEN 1
 					-- with ReprocessBadModule or ReprocessAlternative or ReprocessDBModuleInsertInvalid
 					WHEN status = 540 OR status = 541 OR status = 542 THEN 2
 					ELSE 5
 				END
 			-- non-latest
 			WHEN status = 503 or status = 520 OR status = 521 THEN 3
 			WHEN status = 540 OR status = 541 OR status = 542 THEN 4
 			ELSE 5
 		END,
 		md5(module_path||version) -- deterministic but effectively random
 	LIMIT $2
 `
	// Copyright 2020 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package postgres

	import (
	"context"
	"database/sql"
	"fmt"
	"net/http"
	"strconv"

	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/experiment"
	"golang.org/x/pkgsite/internal/log"
	)

	// UpdateModuleVersionStatesForReprocessing marks modules to be reprocessed
	// that were processed prior to the provided appVersion.
	func (db *DB) UpdateModuleVersionStatesForReprocessing(ctx context.Context, appVersion string) (err error) {
	defer derrors.Wrap(&err, "UpdateModuleVersionStatesForReprocessing(ctx, %q)", appVersion)

	for _, status := range []int{
	http.StatusOK,
	derrors.ToStatus(derrors.HasIncompletePackages),
	derrors.ToStatus(derrors.DBModuleInsertInvalid),
	} {
	if err := db.UpdateModuleVersionStatesWithStatus(ctx, status, appVersion); err != nil {
	return err
	}
	}
	return nil
	}

	func (db *DB) UpdateModuleVersionStatesWithStatus(ctx context.Context, status int, appVersion string) (err error) {
	query := `UPDATE module_version_states
	SET
	status = $2,
	next_processed_after = CURRENT_TIMESTAMP,
	last_processed_at = NULL
	WHERE
	app_version < $1
	AND status = $3;`
	affected, err := db.db.Exec(ctx, query, appVersion, derrors.ToReprocessStatus(status), status)
	if err != nil {
	return err
	}
	log.Infof(ctx,
	"Updated module_version_states with status=%d and app_version < %q to status=%d; %d affected",
	status, appVersion, derrors.ToReprocessStatus(status), affected)
	return nil
	}

	// largeModulePackageThresold represents the package threshold at which it
	// becomes difficult to process packages. Modules with more than this number
	// of packages are generally different versions or forks of kubernetes,
	// aws-sdk-go, azure-sdk-go, and bilibili.
	const largeModulePackageThreshold = 1500

	// largeModulesLimit represents the number of large modules that we are
	// willing to enqueue at a given time.
	// var for testing.
	var largeModulesLimit = config.GetEnvInt("GO_DISCOVERY_LARGE_MODULES_LIMIT", 100)

	// GetNextModulesToFetch returns the next batch of modules that need to be
	// processed. We prioritize modules based on (1) whether it has status zero
	// (never processed), (2) whether it is the latest version, (3) if it is an
	// alternative module, and (4) the number of packages it has. We want to leave
	// time-consuming modules until the end and process them at a slower rate to
	// reduce database load and timeouts. We also want to leave alternative modules
	// towards the end, since these will incur unnecessary deletes otherwise.
	func (db DB) GetNextModulesToFetch(ctx context.Context, limit int) (_ []internal.ModuleVersionState, err error) {
	defer derrors.Wrap(&err, "GetNextModulesToFetch(ctx, %d)", limit)
	queryFmt := nextModulesToProcessQuery
	if experiment.IsActive(ctx, internal.ExperimentAltRequeue) {
	log.Infof(ctx, "using alternative requeue query")
	queryFmt = nextModulesToProcessQueryAlt
	}

	var mvs []*internal.ModuleVersionState
	query := fmt.Sprintf(queryFmt, moduleVersionStateColumns)

	collect := func(rows *sql.Rows) error {
	// Scan the last two columns separately; they are in the query only for sorting.
	scan := func(dests ...interface{}) error {
	var (
	latest bool
	npkg int
	)
	return rows.Scan(append(dests, &latest, &npkg)...)
	}
	mv, err := scanModuleVersionState(scan)
	if err != nil {
	return err
	}
	mvs = append(mvs, mv)
	return nil
	}
	if err := db.db.RunQuery(ctx, query, collect, largeModulePackageThreshold, limit); err != nil {
	return nil, err
	}
	if len(mvs) == 0 {
	log.Infof(ctx, "No modules to requeue")
	} else {
	fmtIntp := func(p *int) string {
	if p == nil {
	return "NULL"
	}
	return strconv.Itoa(*p)
	}
	start := mvs[0]
	end := mvs[len(mvs)-1]
	pkgRange := fmt.Sprintf("%s <= num_packages <= %s", fmtIntp(start.NumPackages), fmtIntp(end.NumPackages))
	log.Infof(ctx, "GetNextModulesToFetch: num_modules=%d; %s; start_module=%q; end_module=%q",
	len(mvs), pkgRange,
	fmt.Sprintf("%s/@v/%s", start.ModulePath, start.Version),
	fmt.Sprintf("%s/@v/%s", end.ModulePath, end.Version))
	}

	// Don't return more than largeModulesLimit of modules that have more than
	// largeModulePackageThreshold packages, or of modules with status zero.
	nLargeOrZero := 0
	for i, m := range mvs {
	if m.Status == 0 \|\| (m.NumPackages != nil && *m.NumPackages >= largeModulePackageThreshold) {
	nLargeOrZero++
	}
	if nLargeOrZero > largeModulesLimit {
	return mvs[:i], nil
	}
	}
	return mvs, nil
	}

	const nextModulesToProcessQuery = `
	-- Make a table of the latest versions of each module.
	WITH latest_versions AS (
	SELECT DISTINCT ON (module_path) module_path, version
	FROM module_version_states
	ORDER BY
	module_path,
	incompatible,
	right(sort_version, 1) = '~' DESC, -- prefer release versions
	sort_version DESC
	)
	SELECT %s, latest, npkg
	FROM (
	SELECT
	%[1]s,
	((module_path, version) IN (SELECT * FROM latest_versions)) AS latest,
	COALESCE(num_packages, 0) AS npkg
	FROM module_version_states
	) s
	WHERE next_processed_after < CURRENT_TIMESTAMP
	AND (status = 0 OR status >= 500)
	ORDER BY
	CASE
	-- new modules
	WHEN status = 0 THEN 0
	WHEN npkg < $1 THEN -- non-large modules
	CASE
	WHEN latest THEN -- latest version
	CASE
	-- with ReprocessStatusOK or ReprocessHasIncompletePackages
	WHEN status = 520 OR status = 521 THEN 1
	-- with ReprocessBadModule or ReprocessAlternative or ReprocessDBModuleInsertInvalid
	WHEN status = 540 OR status = 541 OR status = 542 THEN 2
	ELSE 9
	END
	ELSE -- non-latest version
	CASE
	WHEN status = 520 OR status = 521 THEN 3
	WHEN status = 540 OR status = 541 OR status = 542 THEN 4
	ELSE 9
	END
	END
	ELSE -- large modules
	CASE
	WHEN latest THEN -- latest version
	CASE
	WHEN status = 520 OR status = 521 THEN 5
	WHEN status = 540 OR status = 541 OR status = 542 THEN 6
	ELSE 9
	END
	ELSE -- non-latest version
	CASE
	WHEN status = 520 OR status = 521 THEN 7
	WHEN status = 540 OR status = 541 OR status = 542 THEN 8
	ELSE 9
	END
	END
	END,
	npkg, -- prefer fewer packages
	module_path,
	version -- for reproducibility
	LIMIT $2
	`

	// Like the standard query above, this query prioritizes latest versions.
	// However, it does not hold off large modules until the end. In fact, it tries
	// to avoid grouping modules in any way except by latest and status code:
	// processing is much smoother when they are enqueued in random order.
	//
	// To make the result deterministic for testing, we hash the module path and version
	// rather than actually choosing a random number. md5 is built in to postgres and
	// is an adequate hash for this purpose.
	const nextModulesToProcessQueryAlt = `
	-- Make a table of the latest versions of each module.
	WITH latest_versions AS (
	SELECT DISTINCT ON (module_path) module_path, version
	FROM module_version_states
	ORDER BY
	module_path,
	incompatible,
	right(sort_version, 1) = '~' DESC, -- prefer release versions
	sort_version DESC
	)
	SELECT %s, latest, npkg
	FROM (
	SELECT
	%[1]s,
	((module_path, version) IN (SELECT * FROM latest_versions)) AS latest,
	COALESCE(num_packages, 0) AS npkg
	FROM module_version_states
	) s
	WHERE next_processed_after < CURRENT_TIMESTAMP
	AND (status = 0 OR status >= 500)
	AND $1 = $1 -- ignore largeModulePackageThreshold
	ORDER BY
	CASE
	-- new modules
	WHEN status = 0 THEN 0
	WHEN latest THEN
	CASE
	-- with SheddingLoad or ReprocessStatusOK or ReprocessHasIncompletePackages
	WHEN status = 503 or status = 520 OR status = 521 THEN 1
	-- with ReprocessBadModule or ReprocessAlternative or ReprocessDBModuleInsertInvalid
	WHEN status = 540 OR status = 541 OR status = 542 THEN 2
	ELSE 5
	END
	-- non-latest
	WHEN status = 503 or status = 520 OR status = 521 THEN 3
	WHEN status = 540 OR status = 541 OR status = 542 THEN 4
	ELSE 5
	END,
	md5(module_path\|\|version) -- deterministic but effectively random
	LIMIT $2
	`