init

2019-08-08 17:13:34 +08:00
parent e8e5966a8c
commit f0bf2d5fec
2885 changed files with 1195993 additions and 12 deletions
@@ -0,0 +1,204 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/julienschmidt/httprouter"
+	"github.com/snail007/go-activerecord/mysql"
+)
+
+type UserTraffic struct {
+	ClusterID   uint64
+	UserID      string
+	ServerToken string
+	TunnelID    uint64
+	ConnCount   uint64
+	Upload      uint64
+	Download    uint64
+	Bytes       uint64
+}
+
+var trafficChn = make(chan UserTraffic, 50000)
+
+func init() {
+	go func() {
+		db := G.DB()
+		for {
+			traffic := <-trafficChn
+			if traffic.Bytes > 0 {
+				//update user traffic
+				now := time.Now().Unix()
+				sql := db.AR().Update("package", map[string]interface{}{
+					"bytes_left -": traffic.Bytes,
+					"update_time":  now,
+				}, map[string]interface{}{
+					"user_id":       traffic.UserID,
+					"start_time <=": now,
+					"end_time >=":   now,
+					"bytes_left >":  0,
+				}).OrderBy("end_time", "ASC").Limit(1)
+				_, err := db.Exec(sql)
+				if err != nil {
+					log.Warnf("traffic counter ERR:%s", err)
+				}
+			}
+			if traffic.ConnCount > 0 {
+				//update tunnel conn count
+				rs, err := db.Query(db.AR().From("server").Where(map[string]interface{}{
+					"token": traffic.ServerToken,
+				}).Limit(0, 1))
+				if err != nil {
+					log.Warnf("traffic counter ERR:%s", err)
+				}
+				if rs.Len() == 0 {
+					log.Warnf("traffic counter ERR:empty sever for token:%s", traffic.ServerToken)
+				}
+
+				_, err = db.Exec(db.AR().Replace("conn", map[string]interface{}{
+					"user_id":     traffic.UserID,
+					"cluster_id":  traffic.ClusterID,
+					"server_id":   rs.Value("server_id"),
+					"tunnel_id":   traffic.TunnelID,
+					"count":       traffic.ConnCount,
+					"upload":      traffic.Upload,
+					"download":    traffic.Download,
+					"update_time": time.Now().Unix(),
+				}))
+				if err != nil {
+					log.Warnf("traffic counter ERR:%s", err)
+				}
+			}
+		}
+	}()
+}
+
+type Traffic struct{}
+
+func NewTraffic() *Traffic {
+	return &Traffic{}
+}
+
+//Traffic update a tunnel traffic bytes used,对应cluster的:url.traffic
+//method : POST
+//params : json {"1":{"positive":3231,"negative":5324,"connCount":32."serverToken":"0840d2i30ofs"}}
+func (tunnel *Traffic) Traffic(responseWrite http.ResponseWriter, request *http.Request, params httprouter.Params) {
+	defer request.Body.Close()
+	body, err := ioutil.ReadAll(request.Body)
+	if err != nil {
+		jsonError(responseWrite, err, nil)
+		return
+	}
+	jsonStr := string(body)
+	dataMap := map[string]map[string]interface{}{}
+	err = json.Unmarshal([]byte(jsonStr), &dataMap)
+	if err != nil {
+		log.Warnf("update traffic ERR:%s", err)
+		fmt.Fprint(responseWrite, err.Error())
+		return
+	}
+	if len(dataMap) == 0 {
+		log.Warnf("update traffic ERR:empty data")
+		fmt.Fprint(responseWrite, err.Error())
+		return
+	}
+	tunnelIDs := []string{}
+	for k := range dataMap {
+		tunnelIDs = append(tunnelIDs, k)
+	}
+	db := G.DB()
+	var rs *mysql.ResultSet
+	addr := request.RemoteAddr
+	clusterIP := addr[0:strings.Index(addr, ":")]
+	rs, err = db.Query(db.AR().
+		Select("cluster_id").
+		From("cluster").
+		Where(map[string]interface{}{
+			"ip":        clusterIP,
+			"is_delete": 0,
+		}).Limit(1))
+	if err != nil {
+		log.Warnf("update traffic ERR:%s", err)
+		fmt.Fprint(responseWrite, err.Error())
+		return
+	}
+	if rs.Len() == 0 {
+		err = fmt.Errorf("update traffic ERR:empty cluster for ip :%s", clusterIP)
+		log.Warnf("%s", err)
+		fmt.Fprint(responseWrite, err.Error())
+		return
+	}
+	_clusterID := rs.Value("cluster_id")
+	clusterID, _ := strconv.ParseUint(_clusterID, 10, 64)
+	rs, err = db.Query(db.AR().
+		Select("user_id,tunnel_id").
+		From("tunnel").
+		Where(map[string]interface{}{
+			"tunnel_id": tunnelIDs,
+			"is_delete": 0,
+		}))
+	if err != nil {
+		log.Warnf("update traffic ERR:%s", err)
+		fmt.Fprint(responseWrite, err.Error())
+		return
+	}
+	userIds := rs.MapValues("tunnel_id", "user_id")
+	_data := []UserTraffic{}
+	data := []UserTraffic{}
+	for k, v := range dataMap {
+		userID, ok := userIds[k]
+		if !ok {
+			continue
+		}
+		nVal, err1 := strconv.ParseUint(fmt.Sprintf("%.0f", v["negative"].(float64)), 10, 64)
+		pVal, err2 := strconv.ParseUint(fmt.Sprintf("%.0f", v["positive"].(float64)), 10, 64)
+		connCount, err3 := strconv.ParseUint(fmt.Sprintf("%.0f", v["connCount"].(float64)), 10, 64)
+		interval, err4 := strconv.ParseUint(fmt.Sprintf("%.0f", v["interval"].(float64)), 10, 64)
+		if err1 != nil || err2 != nil || err3 != nil || err4 != nil {
+			log.Warnf("update traffic ERR:%s,%s,%s", err1, err2, err3)
+			fmt.Fprint(responseWrite, "data format incorrect")
+			return
+		}
+		tunnelID, _ := strconv.ParseUint(k, 10, 64)
+		_data = append(_data, UserTraffic{
+			UserID:      userID,
+			Bytes:       uint64(nVal + pVal),
+			ConnCount:   connCount,
+			ServerToken: v["serverToken"].(string),
+			TunnelID:    tunnelID,
+			ClusterID:   clusterID,
+			Upload:      pVal / interval,
+			Download:    nVal / interval,
+		})
+	}
+
+	for _, v := range _data {
+		found := false
+		for k1, v1 := range data {
+			if v.UserID == v1.UserID {
+				found = true
+				v1.Bytes = v.Bytes + v1.Bytes
+				data[k1] = v1
+				break
+			}
+		}
+		if !found {
+			data = append(data, v)
+		}
+	}
+	if len(data) > 0 {
+		for _, v := range data {
+			select {
+			case trafficChn <- v:
+			default:
+				log.Warnf("traffic counter queue is full")
+			}
+		}
+	}
+	responseWrite.WriteHeader(http.StatusNoContent)
+}