From 98ec12266d687b0214c5a221f8038558e9a5d132 Mon Sep 17 00:00:00 2001 From: Matthew R Date: Tue, 7 Jul 2020 21:15:10 -0400 Subject: [PATCH] changes to go func --- Makefile | 2 +- src/go/storage/storage.go | 16 ++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/Makefile b/Makefile index 6bd5ba5..39b0916 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,7 @@ check_cert_signatures: file dist/bin/checkCertSignatures storage: - HOME=/root go build -ldflags="-s -w" -o dist/bin/storage ${storage_files} + HOME=/root go build -ldflags="-s -w" -buildmode=pie -o dist/bin/storage ${storage_files} @chmod 740 dist/bin/storage file dist/bin/storage diff --git a/src/go/storage/storage.go b/src/go/storage/storage.go index 3fdaa61..dabfd4c 100644 --- a/src/go/storage/storage.go +++ b/src/go/storage/storage.go @@ -8,6 +8,7 @@ import ( "os" "strings" "time" + "sync" "github.com/go-redis/redis/v7" "go.mongodb.org/mongo-driver/bson" @@ -81,28 +82,31 @@ func handler(config ConfigStruct) { collection := mongoClient.Database("cloudservices").Collection("accounts") cur, err := collection.Find(context.TODO(), bson.D{}) HandleError(err, 0) - + + var wg sync.WaitGroup for cur.Next(context.TODO()) { - checkAccountSizeAndUpdate(cur.Current.Lookup("username").String(), cur.Current.Lookup("id").String()) + wg.Add(1) + go checkAccountSizeAndUpdate(cur.Current.Lookup("username").String(), cur.Current.Lookup("id").String(), &wg) fmt.Printf("Checking account information for %s\n", cur.Current.Lookup("username").String()) - time.Sleep(300000 * time.Millisecond) } + wg.Wait() err = mongoClient.Disconnect(ctx) HandleError(err, 1) fmt.Println("CSD-GO-STR finished, exiting with code 0. [GO]") os.Exit(0) } -func checkAccountSizeAndUpdate(username string, id string) { +func checkAccountSizeAndUpdate(username string, id string, wg *sync.WaitGroup) { var size float64 = 0 var userHomeDirectory string = strings.Replace(strings.Join([]string{"/home/", string(username)}, ""), "\"", "", -1) usernameFormat := strings.Replace(username, "\"", "", -1) sizeHome := DirSize(&userHomeDirectory) size += sizeHome status := RedisClient.Set("storage"+"-"+usernameFormat, size, 0) - fmt.Println(status.Name()) + fmt.Printf("Redis Status: %s [GO]\n", status.Name()) if status.Err() != nil { fmt.Println(status.Err()) } - fmt.Printf("Set Call | Username: %v, ID: %v | Bytes: %f\n [GO]", string(username), string(id), size) + fmt.Printf("Set Call | Username: %v, ID: %v | Bytes: %f [GO]\n", string(username), string(id), size) + wg.Done() }