From 10c5cf55439fa846fe8bf58d717e49459f1cf437 Mon Sep 17 00:00:00 2001
From: kelvin <kidult.x@gmail.com>
Date: Wed, 19 Jul 2023 21:46:36 +0800
Subject: [PATCH] update

---
 main.go    |  1 +
 service.go | 70 ++++++++++++++++++++++--------------------------------
 2 files changed, 29 insertions(+), 42 deletions(-)

diff --git a/main.go b/main.go
index 997a2eb..c357148 100644
--- a/main.go
+++ b/main.go
@@ -134,6 +134,7 @@ type ModelInfo struct {
 	ModelPrice      float64 `json:"model_price"`
 	ModelPrice2     float64 `json:"model_price2"`
 	ModelPrepayment int     `json:"model_prepayment"`
+	ServerId        int     `json:"server_id"`
 }
 
 var (
diff --git a/service.go b/service.go
index 2683d0e..3cc3467 100644
--- a/service.go
+++ b/service.go
@@ -59,12 +59,24 @@ func queryBlance(key string) (float64, error) {
 // 余额查询
 func checkBlance(key string, model string) (ServerInfo, error) {
 	var serverInfo ServerInfo
+
+	//获取模型价格
+	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
+	if err != nil {
+		return serverInfo, errors.New("模型信息不存在")
+	}
+	var modelInfo ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return serverInfo, errors.New("模型信息解析失败")
+	}
+
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
 	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
+	serverInfoStr, err := Redis.Get(context.Background(), "server:"+string(modelInfo.ServerId)).Result()
 	if err != nil {
 		return serverInfo, errors.New("服务器信息不存在")
 	}
@@ -72,16 +84,7 @@ func checkBlance(key string, model string) (ServerInfo, error) {
 	if err != nil {
 		return serverInfo, errors.New("服务器信息解析失败")
 	}
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return serverInfo, errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return serverInfo, errors.New("模型信息解析失败")
-	}
+
 	//计算余额-先扣除指定金额
 	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment) * modelInfo.ModelPrice)).Result()
 	if err != nil {
@@ -99,12 +102,24 @@ func checkBlance(key string, model string) (ServerInfo, error) {
 // 余额查询 for images
 func checkBlanceForImages(key string, model string, n int) (ServerInfo, error) {
 	var serverInfo ServerInfo
+
+	//获取模型价格
+	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
+	if err != nil {
+		return serverInfo, errors.New("模型信息不存在")
+	}
+	var modelInfo ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return serverInfo, errors.New("模型信息解析失败")
+	}
+
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
 	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
+	serverInfoStr, err := Redis.Get(context.Background(), "server:"+string(modelInfo.ServerId)).Result()
 	if err != nil {
 		return serverInfo, errors.New("服务器信息不存在")
 	}
@@ -112,16 +127,7 @@ func checkBlanceForImages(key string, model string, n int) (ServerInfo, error) {
 	if err != nil {
 		return serverInfo, errors.New("服务器信息解析失败")
 	}
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return serverInfo, errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return serverInfo, errors.New("模型信息解析失败")
-	}
+
 	//计算余额-先扣除指定金额
 	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice)).Result()
 	if err != nil {
@@ -138,20 +144,10 @@ func checkBlanceForImages(key string, model string, n int) (ServerInfo, error) {
 
 // 预扣返还
 func checkBlanceReturn(key string, model string) error {
-	var serverInfo ServerInfo
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
-	if err != nil {
-		return errors.New("服务器信息不存在")
-	}
-	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
-	if err != nil {
-		return errors.New("服务器信息解析失败")
-	}
 	//获取模型价格
 	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
 	if err != nil {
@@ -169,20 +165,10 @@ func checkBlanceReturn(key string, model string) error {
 
 // 预扣返还 for images
 func checkBlanceReturnForImages(key string, model string, n int) error {
-	var serverInfo ServerInfo
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
-	if err != nil {
-		return errors.New("服务器信息不存在")
-	}
-	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
-	if err != nil {
-		return errors.New("服务器信息解析失败")
-	}
 	//获取模型价格
 	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
 	if err != nil {