implement a goroutine pool

turns out b2 doesn't like it if you use too many threads
This commit is contained in:
nyanotech 2024-04-07 03:15:48 -07:00
parent 0bf0144deb
commit 2c5911fd23
Signed by: nyanotech
GPG Key ID: D2D0A9E8F160472B

30
main.go
View File

@ -24,7 +24,7 @@ var secretAccessKey = flag.String("secret-access-key", "", "aws secret key")
var updateExpiresWithin = flag.Int("update-expires-within", 0, "only update objects whose lock expires within this many seconds (default 0)")
var lockFor = flag.Int("lock-for", 90*24*3600, "how many seconds to renew the object lock for (default 90 days)")
var wg sync.WaitGroup
var threadCount = flag.Int("threads", 1024, "how many objects to operate on at a time")
func main() {
flag.Parse()
@ -52,6 +52,17 @@ func main() {
Bucket: bucket,
})
objectQueue := make(chan string)
var wg sync.WaitGroup
for i := 0; i < *threadCount; i++ {
wg.Add(1)
go func() {
defer wg.Done()
queueWorker(svc, objectQueue)
}()
}
for paginator.HasMorePages() {
page, err := paginator.NextPage(context.TODO())
if err != nil {
@ -59,17 +70,24 @@ func main() {
}
for _, item := range page.Contents {
wg.Add(1)
go func(svc *s3.Client, key string) {
defer wg.Done()
checkAndRenewObjectLock(svc, key)
}(svc, *item.Key)
objectQueue <- *item.Key
}
}
close(objectQueue)
wg.Wait()
}
func queueWorker(svc *s3.Client, inQueue chan string) {
for {
object, more := <-inQueue
if !more {
return
}
checkAndRenewObjectLock(svc, object)
}
}
func checkAndRenewObjectLock(svc *s3.Client, object string) {
updateHold := false
if *updateExpiresWithin == 0 {