Compare commits

...

27 Commits

Author SHA1 Message Date
Ahrav Dutta 30008b46a5 import assert. 2022-10-21 08:25:37 -07:00
Ahrav Dutta cf2e56e483 Add mutex. 2022-10-21 08:05:39 -07:00
Ahrav Dutta a076410b53 address comments. 2022-10-21 05:21:29 -07:00
ahrav 98edac9dfb Merge branch 'main' into THOG-712-monotic-progress-tracking 2022-10-21 05:08:40 -07:00
Ahrav Dutta 9beee3cbba fix comments. 2022-10-07 07:53:04 -07:00
Ahrav Dutta 4a3c8e852c fix comments. 2022-10-07 07:52:15 -07:00
Ahrav Dutta 53813884b5 fix comments. 2022-10-07 07:49:59 -07:00
Ahrav Dutta caeef5c5ab fix test. 2022-10-07 07:33:09 -07:00
Ahrav Dutta b55013cad3 remove unused field. 2022-10-07 07:25:11 -07:00
ahrav 9dfd307ff4 Merge branch 'main' into THOG-712-monotic-progress-tracking 2022-10-07 06:03:50 -07:00
Ahrav Dutta 5c1aa826b1 Always complete. 2022-10-07 06:03:34 -07:00
Ahrav Dutta e7313cabf9 Complete scans when we go through all enumerations. 2022-10-07 05:59:22 -07:00
Ahrav Dutta 948f3ef375 Fix invalid method. 2022-10-04 15:55:48 -07:00
ahrav 0d755d82eb Merge branch 'main' into THOG-712-monotic-progress-tracking 2022-10-04 11:33:57 -07:00
ahrav 185a7635f3 Merge branch 'main' into THOG-712-monotic-progress-tracking 2022-09-18 20:58:16 -07:00
Ahrav Dutta d0d521123a Add tests to source progress. 2022-09-18 17:48:27 -07:00
Ahrav Dutta bcab69b8ee Update struct layout. 2022-09-17 18:27:14 -07:00
Ahrav Dutta c84cf53ed6 Update method name to Get. 2022-09-17 18:25:33 -07:00
Ahrav Dutta d6d6940cf9 Include a total counter. This will be used to distinguish total scans regardless of success or not. this is needed in order to display initial progress when starting scans. 2022-09-17 18:22:37 -07:00
Ahrav Dutta e0e29f5b5f Fix Gitlab tests. 2022-09-17 17:43:07 -07:00
Ahrav Dutta aa94718ba4 Add logic to only increment progress counter if a scan succeeds. 2022-09-17 17:30:17 -07:00
Ahrav Dutta 7e84313366 fix tests for github and include tests with new progress logic. 2022-09-17 16:45:34 -07:00
Ahrav Dutta f1a9dcff3f Fix tests for gitlab. 2022-09-17 16:29:04 -07:00
Ahrav Dutta c5662d76a2 Update additional source Progress methods. 2022-09-17 16:20:43 -07:00
Ahrav Dutta 9e8b24ae6b Add extra method to handle scan progress completion for a source. 2022-09-17 16:18:01 -07:00
Ahrav Dutta 8ed8d6ab04 Update commentary for Counter Inc method. 2022-09-17 16:01:06 -07:00
Ahrav Dutta d9f296e7d3 Initial implementation to make progress increment monotically. 2022-09-17 15:59:04 -07:00
10 changed files with 387 additions and 145 deletions
+3 -1
View File
@@ -156,7 +156,9 @@ func main() {
}
for _, repo := range strings.Split(*scanCmdRepo, ",") {
sem.Acquire(ctx, 1)
if err := sem.Acquire(ctx, 1); err != nil {
log.Fatal(err)
}
wgChunkers.Add(1)
go func(r string) {
defer sem.Release(1)
+2 -2
View File
@@ -79,8 +79,8 @@ func (s *Source) Init(aCtx context.Context, name string, jobId, sourceId int64,
// Chunks emits chunks of bytes over a channel.
func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) error {
for i, path := range s.paths {
s.SetProgressComplete(i, len(s.paths), fmt.Sprintf("Path: %s", path), "")
for _, path := range s.paths {
s.Update(len(s.paths), fmt.Sprintf("Path: %s", path), "")
cleanPath := filepath.Clean(path)
done := false
+13 -9
View File
@@ -132,8 +132,8 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
user := cred.BasicAuth.Username
token := cred.BasicAuth.Password
for i, repoURI := range s.conn.Repositories {
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
for _, repoURI := range s.conn.Repositories {
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
if len(repoURI) == 0 {
continue
}
@@ -150,8 +150,8 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
}
}
case *sourcespb.Git_Unauthenticated:
for i, repoURI := range s.conn.Repositories {
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
for _, repoURI := range s.conn.Repositories {
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
if len(repoURI) == 0 {
continue
}
@@ -168,8 +168,8 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
}
}
case *sourcespb.Git_SshAuth:
for i, repoURI := range s.conn.Repositories {
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
for _, repoURI := range s.conn.Repositories {
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
if len(repoURI) == 0 {
continue
}
@@ -189,8 +189,12 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
return errors.New("invalid connection type for git source")
}
for i, u := range s.conn.Directories {
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", u), "")
for _, u := range s.conn.Directories {
s.Counter.IncTotal()
go func() {
defer s.Counter.IncSuccess()
}()
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", u), "")
if len(u) == 0 {
continue
@@ -215,7 +219,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
}
}
s.SetProgressComplete(len(s.conn.Repositories), len(s.conn.Repositories), fmt.Sprintf("Completed scanning source %s", s.name), "")
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
return nil
}
+14 -13
View File
@@ -496,15 +496,17 @@ func (s *Source) scan(ctx context.Context, installationClient *github.Client, ch
s.repos = reposToScan
var scanErrs []error
for i, repoURL := range s.repos {
i, repoURL := i, repoURL
repoCnt := len(s.repos)
for _, repoURL := range s.repos {
repoURL := repoURL
s.jobPool.Go(func() error {
if common.IsDone(ctx) {
return nil
}
cnt := s.Counter.IncTotal()
// TODO: set progress complete is being called concurrently with i
s.setProgressCompleteWithRepo(i, progressIndexOffset, repoURL)
defer s.Counter.IncSuccess()
s.setProgressCompleteWithRepo(progressIndexOffset, repoURL)
// Ensure the repo is removed from the resume info after being scanned.
defer func(s *Source, repoURL string) {
s.resumeInfoMutex.Lock()
@@ -512,15 +514,15 @@ func (s *Source) scan(ctx context.Context, installationClient *github.Client, ch
s.resumeInfoSlice = sources.RemoveRepoFromResumeInfo(s.resumeInfoSlice, repoURL)
}(s, repoURL)
var err error
if !strings.HasSuffix(repoURL, ".git") {
scanErrs = append(scanErrs, fmt.Errorf("repo %s does not end in .git", repoURL))
return nil
}
s.log.WithField("repo", repoURL).Debugf("attempting to clone repo %d/%d", i+1, len(s.repos))
s.log.WithField("repo", repoURL).Debugf("attempting to clone repo #%d of %d", cnt, repoCnt)
var path string
var repo *gogit.Repository
var err error
path, repo, err = s.cloneRepo(ctx, repoURL, installationClient)
if err != nil {
@@ -537,21 +539,20 @@ func (s *Source) scan(ctx context.Context, installationClient *github.Client, ch
git.ScanOptionHeadCommit(s.conn.Head),
)
log.Debugf("Starting to scan repo #%d of %d: %s", cnt, repoCnt, repo)
if err = s.git.ScanRepo(ctx, repo, path, scanOptions, chunksChan); err != nil {
log.WithError(err).Errorf("unable to scan repo, continuing")
return nil
}
atomic.AddUint64(&scanned, 1)
log.Debugf("scanned %d/%d repos", scanned, len(s.repos))
log.Debugf("scanned %d/%d repos", scanned, repoCnt)
return nil
})
}
_ = s.jobPool.Wait()
if len(scanErrs) == 0 {
s.SetProgressComplete(len(s.repos), len(s.repos), "Completed Github scan", "")
}
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
return scanErrs
}
@@ -1006,8 +1007,8 @@ func (s *Source) normalizeRepos(ctx context.Context) {
}
}
// setProgressCompleteWithRepo calls the s.SetProgressComplete after safely setting up the encoded resume info string.
func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL string) {
// setProgressCompleteWithRepo calls the s.Update after safely setting up the encoded resume info string.
func (s *Source) setProgressCompleteWithRepo(offset int, repoURL string) {
s.resumeInfoMutex.Lock()
defer s.resumeInfoMutex.Unlock()
@@ -1018,5 +1019,5 @@ func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL stri
// Make the resume info string from the slice.
encodedResumeInfo := sources.EncodeResumeInfo(s.resumeInfoSlice)
s.SetProgressComplete(index+offset, len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
s.Update(len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
}
+31 -32
View File
@@ -383,7 +383,7 @@ func Test_setProgressCompleteWithRepo_resumeInfo(t *testing.T) {
for _, tt := range tests {
s.resumeInfoSlice = tt.startingResumeInfoSlice
s.setProgressCompleteWithRepo(0, 0, tt.repoURL)
s.setProgressCompleteWithRepo(0, tt.repoURL)
if !reflect.DeepEqual(s.resumeInfoSlice, tt.wantResumeInfoSlice) {
t.Errorf("s.setProgressCompleteWithRepo() got: %v, want: %v", s.resumeInfoSlice, tt.wantResumeInfoSlice)
}
@@ -394,7 +394,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
repos := []string{"a", "b", "c", "d", "e"}
tests := map[string]struct {
repos []string
index int
offset int
wantPercentComplete int64
wantSectionsCompleted int32
@@ -402,7 +401,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
}{
"starting from the beginning, no offset": {
repos: repos,
index: 0,
offset: 0,
wantPercentComplete: 0,
wantSectionsCompleted: 0,
@@ -410,18 +408,16 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
},
"resume from the third, offset 2": {
repos: repos[2:],
index: 0,
offset: 2,
wantPercentComplete: 40,
wantSectionsCompleted: 2,
wantSectionsRemaining: 5,
},
"resume from the third, on last repo, offset 2": {
repos: repos[2:],
index: 2,
offset: 2,
wantPercentComplete: 80,
wantSectionsCompleted: 4,
"resume from the third, on last repo, offset 3": {
repos: repos[3:],
offset: 3,
wantPercentComplete: 60,
wantSectionsCompleted: 3,
wantSectionsRemaining: 5,
},
}
@@ -430,22 +426,29 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
logger.Out = io.Discard
for _, tt := range tests {
s := &Source{
repos: tt.repos,
log: logger.WithField("no", "output"),
}
t.Run("", func(t *testing.T) {
s := &Source{
repos: tt.repos,
log: logger.WithField("no", "output"),
}
s.setProgressCompleteWithRepo(tt.index, tt.offset, "")
gotProgress := s.GetProgress()
if gotProgress.PercentComplete != tt.wantPercentComplete {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
}
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
}
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
}
// Increment the progress counter by number of repos already processed.
for i := 1; i <= tt.offset; i++ {
s.Counter.IncSuccess()
}
s.setProgressCompleteWithRepo(tt.offset, "")
gotProgress := s.GetProgress()
if gotProgress.PercentComplete != tt.wantPercentComplete {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
}
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
t.Errorf("s.setProgressCompleteWithRepo() SectionsCompleted got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
}
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
t.Errorf("s.setProgressCompleteWithRepo() SectionsRemaining got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
}
})
}
}
@@ -454,16 +457,15 @@ func Test_scan_SetProgressComplete(t *testing.T) {
name string
repos []string
wantComplete bool
wantErr bool
}{
{
name: "no repos",
wantComplete: true,
},
{
name: "one valid repo",
repos: []string{"a"},
wantErr: true,
name: "one valid repo",
repos: []string{"a"},
wantComplete: true,
},
}
@@ -475,9 +477,6 @@ func Test_scan_SetProgressComplete(t *testing.T) {
src.jobPool = &errgroup.Group{}
_ = src.scan(context.Background(), nil, nil)
if !tc.wantErr {
assert.Equal(t, "", src.GetProgress().EncodedResumeInfo)
}
gotComplete := src.GetProgress().PercentComplete == 100
if gotComplete != tc.wantComplete {
+48 -41
View File
@@ -8,6 +8,9 @@ import (
"sort"
"strings"
"sync"
"sync/atomic"
"golang.org/x/sync/errgroup"
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
@@ -46,6 +49,7 @@ type Source struct {
resumeInfoMutex sync.Mutex
jobSem *semaphore.Weighted
sources.Progress
jobPool *errgroup.Group
}
// Ensure the Source satisfies the interface at compile time.
@@ -73,7 +77,8 @@ func (s *Source) Init(aCtx context.Context, name string, jobId, sourceId int64,
s.sourceId = sourceId
s.jobId = jobId
s.verify = verify
s.jobSem = semaphore.NewWeighted(int64(concurrency))
s.jobPool = &errgroup.Group{}
s.jobPool.SetLimit(concurrency)
var conn sourcespb.GitLab
err := anypb.UnmarshalTo(connection, &conn, proto.UnmarshalOptions{})
@@ -269,34 +274,32 @@ func (s *Source) getRepos() ([]string, []error) {
}
func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk) []error {
wg := sync.WaitGroup{}
var errs []error
var errsMut sync.Mutex
var scanned uint64
// If there is resume information available, limit this scan to only the repos that still need scanning.
reposToScan, progressIndexOffset := sources.FilterReposToResume(s.repos, s.GetProgress().EncodedResumeInfo)
s.repos = reposToScan
log.Debugf("Found %v total repos to scan", len(s.repos))
for i, repo := range s.repos {
if common.IsDone(ctx) {
// We are returning nil instead of the scanErrors slice here because
// we don't want to mark this scan as errored if we cancelled it.
return nil
}
if err := s.jobSem.Acquire(ctx, 1); err != nil {
log.WithError(err).Debug("could not acquire semaphore")
continue
}
wg.Add(1)
go func(ctx context.Context, repoURL string, i int) {
defer s.jobSem.Release(1)
defer wg.Done()
var errMut sync.Mutex
var scanErrs []error
repoCnt := len(s.repos)
for _, repoURL := range s.repos {
repoURL := repoURL
s.jobPool.Go(func() error {
if common.IsDone(ctx) {
// We are returning nil instead of the scanErrors slice here because
// we don't want to mark this scan as errored if we cancelled it.
return nil
}
if len(repoURL) == 0 {
return
return nil
}
s.setProgressCompleteWithRepo(i, progressIndexOffset, repoURL)
// Ensure the repo is removed from the resume info after being scanned.
cnt := s.Counter.IncTotal()
defer s.Counter.IncSuccess()
s.setProgressCompleteWithRepo(progressIndexOffset, repoURL)
// Ensure the repoURL is removed from the resume info after being scanned.
defer func(s *Source) {
s.resumeInfoMutex.Lock()
defer s.resumeInfoMutex.Unlock()
@@ -309,7 +312,7 @@ func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk)
if s.authMethod == "UNAUTHENTICATED" {
path, repo, err = git.CloneRepoUsingUnauthenticated(repoURL)
} else {
// If a username is not provided we need to use a default one in order to clone a private repo.
// If a username is not provided we need to use a default one in order to clone a private repoURL.
// Not setting "placeholder" as s.user on purpose in case any downstream services rely on a "" value for s.user.
user := s.user
if user == "" {
@@ -319,28 +322,32 @@ func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk)
}
defer os.RemoveAll(path)
if err != nil {
errsMut.Lock()
errs = append(errs, err)
errsMut.Unlock()
return
errMut.Lock()
scanErrs = append(scanErrs, err)
errMut.Unlock()
return nil
}
log.Debugf("Starting to scan repo %d/%d: %s", i+1, len(s.repos), repoURL)
log.Debugf("Starting to scan repoURL #%d of %d: %s", cnt, repoCnt, repo)
err = s.git.ScanRepo(ctx, repo, path, git.NewScanOptions(), chunksChan)
if err != nil {
errsMut.Lock()
errs = append(errs, err)
errsMut.Unlock()
return
errMut.Lock()
scanErrs = append(scanErrs, err)
errMut.Unlock()
return nil
}
log.Debugf("Completed scanning repo %d/%d: %s", i+1, len(s.repos), repoURL)
}(ctx, repo, i)
}
wg.Wait()
if len(errs) == 0 {
s.SetProgressComplete(len(s.repos), len(s.repos), fmt.Sprintf("Completed scanning source %s", s.name), "")
atomic.AddUint64(&scanned, 1)
log.Debugf("scanned %d/%d repos", scanned, repoCnt)
log.Debugf("Completed scanning repoURL # %d of %d: %s", cnt, repoCnt, repo)
return nil
})
}
return errs
_ = s.jobPool.Wait()
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
return scanErrs
}
// Chunks emits chunks of bytes over a channel.
@@ -417,8 +424,8 @@ func (s *Source) basicAuthSuccessful(apiClient *gitlab.Client) bool {
return false
}
// setProgressCompleteWithRepo calls the s.SetProgressComplete after safely setting up the encoded resume info string.
func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL string) {
// setProgressCompleteWithRepo calls the s.Update after safely setting up the encoded resume info string.
func (s *Source) setProgressCompleteWithRepo(offset int, repoURL string) {
s.resumeInfoMutex.Lock()
defer s.resumeInfoMutex.Unlock()
@@ -430,5 +437,5 @@ func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL stri
encodedResumeInfo := sources.EncodeResumeInfo(s.resumeInfoSlice)
// Add the offset to both the index and the repos to give the proper place and proper repo count.
s.SetProgressComplete(index+offset, len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
s.Update(len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
}
+31 -32
View File
@@ -9,7 +9,7 @@ import (
"github.com/kylelemons/godebug/pretty"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"golang.org/x/sync/semaphore"
"golang.org/x/sync/errgroup"
"google.golang.org/protobuf/types/known/anypb"
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
@@ -187,7 +187,7 @@ func Test_setProgressCompleteWithRepo_resumeInfo(t *testing.T) {
for _, tt := range tests {
s.resumeInfoSlice = tt.startingResumeInfoSlice
s.setProgressCompleteWithRepo(0, 0, tt.repoURL)
s.setProgressCompleteWithRepo(0, tt.repoURL)
if !reflect.DeepEqual(s.resumeInfoSlice, tt.wantResumeInfoSlice) {
t.Errorf("s.setProgressCompleteWithRepo() got: %v, want: %v", s.resumeInfoSlice, tt.wantResumeInfoSlice)
}
@@ -198,7 +198,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
repos := []string{"a", "b", "c", "d", "e"}
tests := map[string]struct {
repos []string
index int
offset int
wantPercentComplete int64
wantSectionsCompleted int32
@@ -206,7 +205,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
}{
"starting from the beginning, no offset": {
repos: repos,
index: 0,
offset: 0,
wantPercentComplete: 0,
wantSectionsCompleted: 0,
@@ -214,18 +212,16 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
},
"resume from the third, offset 2": {
repos: repos[2:],
index: 0,
offset: 2,
wantPercentComplete: 40,
wantSectionsCompleted: 2,
wantSectionsRemaining: 5,
},
"resume from the third, on last repo, offset 2": {
repos: repos[2:],
index: 2,
offset: 2,
wantPercentComplete: 80,
wantSectionsCompleted: 4,
repos: repos[3:],
offset: 3,
wantPercentComplete: 60,
wantSectionsCompleted: 3,
wantSectionsRemaining: 5,
},
}
@@ -233,21 +229,28 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
log.SetOutput(io.Discard)
for _, tt := range tests {
s := &Source{
repos: tt.repos,
}
t.Run("", func(t *testing.T) {
s := &Source{
repos: tt.repos,
}
s.setProgressCompleteWithRepo(tt.index, tt.offset, "")
gotProgress := s.GetProgress()
if gotProgress.PercentComplete != tt.wantPercentComplete {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
}
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
}
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
}
// Increment the progress counter by number of repos already processed.
for i := 1; i <= tt.offset; i++ {
s.Counter.IncSuccess()
}
s.setProgressCompleteWithRepo(tt.offset, "")
gotProgress := s.GetProgress()
if gotProgress.PercentComplete != tt.wantPercentComplete {
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
}
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
t.Errorf("s.setProgressCompleteWithRepo() SectionsCompleted got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
}
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
t.Errorf("s.setProgressCompleteWithRepo() SectionsRemaining got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
}
})
}
}
@@ -256,16 +259,15 @@ func Test_scanRepos_SetProgressComplete(t *testing.T) {
name string
repos []string
wantComplete bool
wantErr bool
}{
{
name: "no repos",
wantComplete: true,
},
{
name: "one valid repo",
repos: []string{"a"},
wantErr: true,
name: "one valid repo",
repos: []string{"a"},
wantComplete: true,
},
}
@@ -274,12 +276,9 @@ func Test_scanRepos_SetProgressComplete(t *testing.T) {
src := &Source{
repos: tc.repos,
}
src.jobSem = semaphore.NewWeighted(1)
src.jobPool = &errgroup.Group{}
_ = src.scanRepos(context.Background(), nil)
if !tc.wantErr {
assert.Equal(t, "", src.GetProgress().EncodedResumeInfo)
}
gotComplete := src.GetProgress().PercentComplete == 100
if gotComplete != tc.wantComplete {
+8 -3
View File
@@ -135,12 +135,15 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
return errors.Errorf("invalid configuration given for %s source", s.name)
}
for i, bucket := range bucketsToScan {
numBuckets := len(bucketsToScan)
log.Debugf("Found %v total buckets to scan", numBuckets)
for _, bucket := range bucketsToScan {
if common.IsDone(ctx) {
return nil
}
s.SetProgressComplete(i, len(bucketsToScan), fmt.Sprintf("Bucket: %s", bucket), "")
cnt := s.Counter.IncSuccess()
s.Update(len(bucketsToScan), fmt.Sprintf("Bucket: %s", bucket), "")
s.log.Debugf("Scanning bucket: %s", bucket)
region, err := s3manager.GetBucketRegionWithClient(context.Background(), client, bucket)
@@ -161,6 +164,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
// pf := "public"
errorCount := sync.Map{}
log.Debugf("Starting to scan bucket #%d of %d: %s", cnt, numBuckets, bucket)
err = regionalClient.ListObjectsV2PagesWithContext(
ctx, &s3.ListObjectsV2Input{Bucket: &bucket},
func(page *s3.ListObjectsV2Output, last bool) bool {
@@ -172,8 +176,9 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
s.log.WithError(err).Errorf("could not list objects in s3 bucket: %s", bucket)
return errors.WrapPrefix(err, fmt.Sprintf("could not list objects in s3 bucket: %s", bucket), 0)
}
log.Debugf("scanned %d/%d buckets", cnt, numBuckets)
}
s.SetProgressComplete(len(bucketsToScan), len(bucketsToScan), fmt.Sprintf("Completed scanning source %s", s.name), "")
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
return nil
}
+45 -12
View File
@@ -2,6 +2,7 @@ package sources
import (
"sync"
"sync/atomic"
"google.golang.org/protobuf/types/known/anypb"
@@ -104,39 +105,71 @@ func NewConfig(opts ...func(*Config)) Config {
return *c
}
// Counter is a simple monotonically-increasing counter for tracking progress.
// It includes a successCnt that records the number of successful scans.
// It also includes a totalCnt that records the total number of scans.
// This is necessary in order to display the number of concurrent scans
// regardless of whether they are successful or not.
type Counter struct {
successCnt, totalCnt uint32
}
// IncTotal increments the total field of counter and returns the new value.
func (c *Counter) IncTotal() int {
return int(atomic.AddUint32(&c.totalCnt, 1))
}
// IncSuccess increments the success field of counter and returns the new value.
func (c *Counter) IncSuccess() int {
return int(atomic.AddUint32(&c.successCnt, 1))
}
// Get returns the current successCnt.
func (c *Counter) Get() uint32 {
return atomic.LoadUint32(&c.successCnt)
}
// Progress is used to update job completion progress across sources.
type Progress struct {
mut sync.Mutex
SectionsCompleted int32
SectionsRemaining int32
PercentComplete int64
Message string
EncodedResumeInfo string
SectionsCompleted int32
SectionsRemaining int32
mut sync.Mutex
Counter Counter
}
// SetProgressComplete sets job progress information for a running job based on the highest level objects in the source.
// i is the current iteration in the loop of target scope
// Update sets job progress information for a running job based on the highest level objects in the source.
// scope should be the len(scopedItems)
// message is the public facing user information about the current progress
// encodedResumeInfo is an optional string representing any information necessary to resume the job if interrupted
func (p *Progress) SetProgressComplete(i, scope int, message, encodedResumeInfo string) {
func (p *Progress) Update(scope int, message, encodedResumeInfo string) {
p.mut.Lock()
defer p.mut.Unlock()
i := p.Counter.Get()
p.Message = message
p.EncodedResumeInfo = encodedResumeInfo
p.SectionsCompleted = int32(i)
p.SectionsRemaining = int32(scope)
// If the iteration and scope are both 0, completion is 100%.
if i == 0 && scope == 0 {
p.PercentComplete = 100
return
}
p.PercentComplete = int64((float64(i) / float64(scope)) * 100)
}
// Complete sets job progress information for a completed job.
func (p *Progress) Complete(msg string) {
p.mut.Lock()
defer p.mut.Unlock()
p.SectionsRemaining = 0
p.SectionsCompleted = int32(p.Counter.Get())
p.Message = msg
p.EncodedResumeInfo = ""
p.PercentComplete = 100
}
// GetProgress gets job completion percentage for metrics reporting.
func (p *Progress) GetProgress() *Progress {
p.mut.Lock()
+192
View File
@@ -0,0 +1,192 @@
package sources
import (
"sync"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"google.golang.org/protobuf/testing/protocmp"
)
func TestCounter_IncTotal(t *testing.T) {
tests := []struct {
name string
want int
}{
{
name: "increment total 50 times",
want: 50,
},
{
name: "increment total 0 times",
},
{
name: "increment total 1 times",
want: 1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
c := &Counter{}
var wg sync.WaitGroup
wg.Add(tt.want)
for i := 0; i < tt.want; i++ {
go func(c *Counter) {
defer wg.Done()
c.IncTotal()
}(c)
}
wg.Wait()
if got := c.totalCnt; int(got) != tt.want {
t.Errorf("Counter.IncTotal() = %v, want %v", got, tt.want)
}
})
}
}
func TestCounter_IncSuccess(t *testing.T) {
tests := []struct {
name string
want int
}{
{
name: "increment total 50 times",
want: 50,
},
{
name: "increment total 0 times",
},
{
name: "increment total 1 times",
want: 1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
c := &Counter{}
var wg sync.WaitGroup
wg.Add(tt.want)
for i := 0; i < tt.want; i++ {
go func(c *Counter) {
defer wg.Done()
c.IncSuccess()
}(c)
}
wg.Wait()
if got := c.Get(); int(got) != tt.want {
t.Errorf("Counter.IncTotal() = %v, want %v", got, tt.want)
}
})
}
}
func TestProgress_Update(t *testing.T) {
tests := []struct {
name, msg, encodeResumeInfo string
counter uint32
scope, routines, want int
}{
{
name: "update progress, no counter",
scope: 100,
routines: 10,
want: 10,
},
{
name: "update progress, counter 5",
scope: 100,
routines: 10,
counter: 5,
want: 15,
},
{
name: "update progress, counter 50",
scope: 100,
routines: 10,
counter: 50,
want: 60,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
p := &Progress{}
p.Counter.successCnt = tt.counter
var wg sync.WaitGroup
wg.Add(tt.routines)
for i := 0; i < tt.routines; i++ {
p.Counter.IncSuccess()
go func(p *Progress) {
defer wg.Done()
p.Update(tt.scope, tt.msg, tt.encodeResumeInfo)
}(p)
}
wg.Wait()
if got := p.PercentComplete; int(got) != tt.want {
t.Errorf("Progress.Update(%v, %v, %v): got %v, want %v", tt.scope, tt.msg, tt.encodeResumeInfo, got, tt.want)
}
})
}
}
func TestProgress_Complete(t *testing.T) {
tests := []struct {
name,
msg string
curProgress, wantProgress *Progress
}{
{
name: "complete progress",
msg: "completed scanning thinger",
curProgress: &Progress{
PercentComplete: 42,
Message: "scanning thinger",
EncodedResumeInfo: "resume info",
SectionsRemaining: 12,
SectionsCompleted: 2,
},
wantProgress: &Progress{
PercentComplete: 100,
Message: "completed scanning thinger",
EncodedResumeInfo: "",
SectionsRemaining: 0,
SectionsCompleted: 14,
},
},
{
name: "complete progress with no message",
curProgress: &Progress{
PercentComplete: 80,
Message: "scanning thinger2",
EncodedResumeInfo: "resume info2",
SectionsRemaining: 2,
SectionsCompleted: 42,
},
wantProgress: &Progress{
PercentComplete: 100,
EncodedResumeInfo: "",
SectionsRemaining: 0,
SectionsCompleted: 44,
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tt.curProgress.Counter.successCnt = uint32(tt.wantProgress.SectionsCompleted)
tt.curProgress.Complete(tt.msg)
ignoreOpts := cmpopts.IgnoreFields(Progress{}, "mut", "Counter")
if !cmp.Equal(tt.curProgress, tt.wantProgress, protocmp.Transform(), ignoreOpts) {
t.Errorf("Progress.Complete(%v): got %v, want %v", tt.msg, tt.curProgress, tt.wantProgress)
}
})
}
}