anytunnel/at-api/traffic.go

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)
}