test/ping/utils/utils.go

package utils

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/wecisecode/util/logger"
)

var (
	dataStats   = make(map[string]*DataStat) // <title-name>:DataStat
	dataStatsMu sync.Mutex
)

func CheckWritePermission(fp string) (err error) {
	var file *os.File
	if file, err = ioutil.TempFile(fp, ".toe.perm_check"); err == nil {
		errMsgs := make([]string, 0, 3)
		var e error
		if _, e = file.WriteString("ok"); e != nil {
			errMsgs = append(errMsgs, "can't write to test file")
		}
		if e = file.Close(); e != nil {
			errMsgs = append(errMsgs, "can't close test file")
		}
		if e = os.Remove(file.Name()); e != nil {
			errMsgs = append(errMsgs, "can't remove test file")
		}
		if len(errMsgs) > 0 {
			err = fmt.Errorf("errors: %s", strings.Join(errMsgs, ", "))
		}
	}
	return
}

type DataStat struct {
	statChan     chan int
	running      bool
	title        string
	name         string
	Interval     int   `json:"interval"`
	Success      int64 `json:"success"`
	Fail         int64 `json:"fail"`
	Ignore       int64 `json:"ignore"`
	TotalSuccess int64 `json:"total_success"`
	TotalFail    int64 `json:"total_fail"`
	TotalIgnore  int64 `json:"total_ignore"`
	StartTime    int64 `json:"start_time"`
	LastTime     int64 `json:"last_time"`
	StatTime     int64 `json:"stat_time"`
}

func NewDataStat(title, name string, interval int, chanSize int) *DataStat {
	dataStatsMu.Lock()
	defer dataStatsMu.Unlock()
	key := title + "-" + name
	if ds, ok := dataStats[key]; ok {
		if ds.statChan != nil {
			close(ds.statChan)
		}
		delete(dataStats, key)
	}

	ds := &DataStat{
		title:        title,
		name:         name,
		statChan:     make(chan int, chanSize),
		Interval:     interval,
		Success:      0,
		Fail:         0,
		TotalSuccess: 0,
		TotalFail:    0,
	}
	dataStats[key] = ds

	return ds
}

func GetDataStats() []*DataStat {
	dataStatsMu.Lock()
	defer dataStatsMu.Unlock()
	var list []*DataStat
	for _, ds := range dataStats {
		list = append(list, ds)
	}
	return list
}

func GetDataStat(key string) *DataStat {
	dataStatsMu.Lock()
	defer dataStatsMu.Unlock()
	return dataStats[key]
}

func (ds *DataStat) Close() {
	if ds.statChan != nil && ds.running {
		close(ds.statChan)
	}
	ds.statChan = nil
}

func (ds *DataStat) Start() {
	tc := time.NewTicker(time.Duration(ds.Interval) * time.Second)
	startTime := time.Now()
	ds.StartTime = startTime.UnixNano() / int64(time.Millisecond)
	ds.running = true
	defer func() {
		ds.running = false
	}()
L:
	for {
		select {
		case <-tc.C:
			now := time.Now()
			speed := float64(ds.Success+ds.Fail+ds.Ignore) / now.Sub(startTime).Seconds()
			startTime = now
			logger.Infof("Stat %s %s: success:%d fail:%d ignore:%d within %ds, total success:%d total fail:%d total ignore:%d speed:%.1f/s.\n", ds.title, ds.name, ds.Success, ds.Fail, ds.Ignore, ds.Interval, ds.TotalSuccess, ds.TotalFail, ds.TotalIgnore, speed)
			ds.StatTime = now.UnixNano() / int64(time.Millisecond)
			ds.Success = 0
			ds.Fail = 0
			ds.Ignore = 0
		case n, ok := <-ds.statChan:
			if !ok {
				break L
			}
			ds.LastTime = time.Now().UnixNano() / int64(time.Millisecond)
			switch n {
			case 0:
				ds.Success++
				ds.TotalSuccess++
			case 1:
				ds.Fail++
				ds.TotalFail++
			case 2:
				ds.Ignore++
				ds.TotalIgnore++
			}
		}
	}
}

func (ds *DataStat) Send(status int) {
	if ds.statChan != nil && ds.running {
		ds.statChan <- status
	}
}

func (ds *DataStat) GetTitle() string {
	return ds.title
}

func (ds *DataStat) GetName() string {
	return ds.name
}

//func Stats(title, name string, interval int, statChan chan int) {
//	var totalSuccess, success, fail, totalFail int64
//	tc := time.NewTicker(time.Duration(interval) * time.Second)
//	lastTime := time.Now()
//L:
//	for {
//		select {
//		case <-tc.C:
//			now := time.Now()
//			speed := float64(success + fail) / now.Sub(lastTime).Seconds()
//			lastTime = now
//			logger.Infof("Stats %s %s: success:%d fail:%d within %ds, total success:%d total fail:%d speed:%.1f/s.\n", title, name, success, fail, interval, totalSuccess, totalFail, speed)
//			success = 0
//			fail = 0
//			//ignore = 0
//		case n, ok := <-statChan:
//			if !ok {
//				break L
//			}
//			switch n {
//			case 0:
//				success++
//				totalSuccess++
//			case 1:
//				fail++
//				totalFail++
//				//case 2:
//				//	ignore ++
//				//	totalIgnore ++
//			}
//		}
//	}
//}