using channels and semaphore logic to do archivation concurrently

weni-ai · Jan 7, 2025 · a81f151 · a81f151
1 parent 6f65fa3
commit a81f151
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 18 deletions.
diff --git a/archives/config.go b/archives/config.go
@@ -28,6 +28,8 @@ type Config struct {
 
 	RollupOrgTimeout          int `help:"rollup timeout for all org archives, limit in hours (default 3)"`
 	BuildRollupArchiveTimeout int `help:"rollup for single archive timeout, limit in hours (default 1)"`
+
+	MaxConcurrentArchivation int `help:"max concurrent org archivation (default 2)"`
 }
 
 // NewConfig returns a new default configuration object
@@ -58,6 +60,8 @@ func NewConfig() *Config {
 
 		RollupOrgTimeout:          3,
 		BuildRollupArchiveTimeout: 1,
+
+		MaxConcurrentArchivation: 2,
 	}
 
 	return &config

diff --git a/cmd/rp-archiver/main.go b/cmd/rp-archiver/main.go
@@ -79,6 +79,29 @@ func main() {
 		logrus.WithError(err).Fatal("cannot write to temp directory")
 	}
 
+	semaphore := make(chan struct{}, config.MaxConcurrentArchivation)
+
+	archiveTask := func(org archives.Org) {
+		// no single org should take more than 12 hours
+		ctx, cancel := context.WithTimeout(context.Background(), time.Hour*12)
+
+		log := logrus.WithField("org", org.Name).WithField("org_id", org.ID)
+
+		if config.ArchiveMessages {
+			_, _, err = archives.ArchiveOrg(ctx, time.Now(), config, db, s3Client, org, archives.MessageType)
+			if err != nil {
+				log.WithError(err).WithField("archive_type", archives.MessageType).Error("error archiving org messages")
+			}
+		}
+		if config.ArchiveRuns {
+			_, _, err = archives.ArchiveOrg(ctx, time.Now(), config, db, s3Client, org, archives.RunType)
+			if err != nil {
+				log.WithError(err).WithField("archive_type", archives.RunType).Error("error archiving org runs")
+			}
+		}
+		cancel()
+	}
+
 	for {
 		start := time.Now().In(time.UTC)
 
@@ -110,24 +133,8 @@ func main() {
 
 		// for each org, do our export
 		for _, org := range orgs {
-			// no single org should take more than 12 hours
-			ctx, cancel := context.WithTimeout(context.Background(), time.Hour*12)
-			log := logrus.WithField("org", org.Name).WithField("org_id", org.ID)
-
-			if config.ArchiveMessages {
-				_, _, err = archives.ArchiveOrg(ctx, time.Now(), config, db, s3Client, org, archives.MessageType)
-				if err != nil {
-					log.WithError(err).WithField("archive_type", archives.MessageType).Error("error archiving org messages")
-				}
-			}
-			if config.ArchiveRuns {
-				_, _, err = archives.ArchiveOrg(ctx, time.Now(), config, db, s3Client, org, archives.RunType)
-				if err != nil {
-					log.WithError(err).WithField("archive_type", archives.RunType).Error("error archiving org runs")
-				}
-			}
-
-			cancel()
+			semaphore <- struct{}{}
+			go archiveTask(org)
 		}
 
 		// ok, we did all our work for our orgs, quit if so configured or sleep until the next day