cloud: split shepherd up
Change-Id: I8e386d9eaaf17543743e1e8a37a8d71426910d59
Reviewed-on: https://review.monogon.dev/c/monogon/+/2213
Reviewed-by: Serge Bazanski <serge@monogon.tech>
Tested-by: Jenkins CI
diff --git a/go/mflags/limiter.go b/go/mflags/limiter.go
new file mode 100644
index 0000000..c085b52
--- /dev/null
+++ b/go/mflags/limiter.go
@@ -0,0 +1,40 @@
+// Package mflags implements custom flags for use in monogon projects.
+// It provides them to be used like normal flag.$Var and registers the
+// required flag functions.
+package mflags
+
+import (
+ "flag"
+ "fmt"
+ "strconv"
+ "strings"
+ "time"
+
+ "golang.org/x/time/rate"
+)
+
+// Limiter configures a *rate.Limiter as a flag.
+func Limiter(l **rate.Limiter, name, defval, help string) {
+ syntax := "'duration,count' eg. '2m,10' for a 10-sized bucket refilled at one token every 2 minutes"
+ help = help + fmt.Sprintf(" (default: %q, syntax: %s)", defval, syntax)
+ flag.Func(name, help, func(val string) error {
+ if val == "" {
+ val = defval
+ }
+ parts := strings.Split(val, ",")
+ if len(parts) != 2 {
+ return fmt.Errorf("invalid syntax, want: %s", syntax)
+ }
+ duration, err := time.ParseDuration(parts[0])
+ if err != nil {
+ return fmt.Errorf("invalid duration: %w", err)
+ }
+ refill, err := strconv.ParseUint(parts[1], 10, 31)
+ if err != nil {
+ return fmt.Errorf("invalid refill rate: %w", err)
+ }
+ *l = rate.NewLimiter(rate.Every(duration), int(refill))
+ return nil
+ })
+ flag.Set(name, defval)
+}