调整个限流参数

icowan · Feb 25, 2020 · ec0fffa · ec0fffa
1 parent ba19b02
commit ec0fffa
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/cmd/service/service.go b/cmd/service/service.go
@@ -63,6 +63,7 @@ func Run() {
 	logPath = envString("LOG_PATH", logPath)
 	logLevel = envString("LOG_LEVEL", logLevel)
 	devCors = envString("DEV_CORS", devCors)
+	rateBucketNum = envInt("RATE_BUCKET", rateBucketNum)
 
 	logger = logging.SetLogging(logger, logPath, logLevel)
 
@@ -168,3 +169,13 @@ func envString(env string, fallback *string) *string {
 	}
 	return &e
 }
+
+func envInt(env string, fallback *int) *int {
+	e := os.Getenv(env)
+	if e == "" {
+		_ = os.Setenv(env, strconv.Itoa(*fallback))
+		return fallback
+	}
+	num, _ := strconv.Atoi(e)
+	return &num
+}