<commit_before>package v1

import (
	"fmt"

	"git.zxq.co/ripple/ocl"
	"git.zxq.co/ripple/rippleapi/common"
)

type leaderboardUser struct {
	userData
	ChosenMode    modeData `json:"chosen_mode"`
	PlayStyle     int      `json:"play_style"`
	FavouriteMode int      `json:"favourite_mode"`
}

type leaderboardResponse struct {
	common.ResponseBase
	Users []leaderboardUser `json:"users"`
}

const lbUserQuery = `
SELECT
	users.id, users.username, users.register_datetime, users.privileges, users.latest_activity,

	users_stats.username_aka, users_stats.country,
	users_stats.play_style, users_stats.favourite_mode,

	users_stats.ranked_score_%[1]s, users_stats.total_score_%[1]s, users_stats.playcount_%[1]s,
	users_stats.replays_watched_%[1]s, users_stats.total_hits_%[1]s,
	users_stats.avg_accuracy_%[1]s, users_stats.pp_%[1]s, leaderboard_%[1]s.position as %[1]s_position
FROM leaderboard_%[1]s
INNER JOIN users ON users.id = leaderboard_%[1]s.user
INNER JOIN users_stats ON users_stats.id = leaderboard_%[1]s.user
%[2]s`

// LeaderboardGET gets the leaderboard.
func LeaderboardGET(md common.MethodData) common.CodeMessager {
	m := getMode(md.Query("mode"))
	// Admins may not want to see banned users on the leaderboard.
	// This is the default setting. In case they do, they have to activate see_everything.
	query := fmt.Sprintf(lbUserQuery, m, `WHERE `+md.User.OnlyUserPublic(md.HasQuery("see_everything"))+
		` ORDER BY leaderboard_`+m+`.position `+common.Paginate(md.Query("p"), md.Query("l"), 500))
	rows, err := md.DB.Query(query)
	if err != nil {
		md.Err(err)
		return Err500
	}
	var resp leaderboardResponse
	for rows.Next() {
		var u leaderboardUser
		err := rows.Scan(
			&u.ID, &u.Username, &u.RegisteredOn, &u.Privileges, &u.LatestActivity,

			&u.UsernameAKA, &u.Country, &u.PlayStyle, &u.FavouriteMode,

			&u.ChosenMode.RankedScore, &u.ChosenMode.TotalScore, &u.ChosenMode.PlayCount,
			&u.ChosenMode.ReplaysWatched, &u.ChosenMode.TotalHits,
			&u.ChosenMode.Accuracy, &u.ChosenMode.PP, &u.ChosenMode.GlobalLeaderboardRank,
		)
		if err != nil {
			md.Err(err)
			continue
		}
		u.ChosenMode.Level = ocl.GetLevelPrecise(int64(u.ChosenMode.TotalScore))
		resp.Users = append(resp.Users, u)
	}
	resp.Code = 200
	return resp
}
<commit_msg>Add ability to filter leaderboard by country (country ranking) (NOT ON WEBSITE RN)<commit_after>package v1

import (
	"fmt"

	"git.zxq.co/ripple/ocl"
	"git.zxq.co/ripple/rippleapi/common"
)

type leaderboardUser struct {
	userData
	ChosenMode    modeData `json:"chosen_mode"`
	PlayStyle     int      `json:"play_style"`
	FavouriteMode int      `json:"favourite_mode"`
}

type leaderboardResponse struct {
	common.ResponseBase
	Users []leaderboardUser `json:"users"`
}

const lbUserQuery = `
SELECT
	users.id, users.username, users.register_datetime, users.privileges, users.latest_activity,

	users_stats.username_aka, users_stats.country,
	users_stats.play_style, users_stats.favourite_mode,

	users_stats.ranked_score_%[1]s, users_stats.total_score_%[1]s, users_stats.playcount_%[1]s,
	users_stats.replays_watched_%[1]s, users_stats.total_hits_%[1]s,
	users_stats.avg_accuracy_%[1]s, users_stats.pp_%[1]s, leaderboard_%[1]s.position as %[1]s_position
FROM leaderboard_%[1]s
INNER JOIN users ON users.id = leaderboard_%[1]s.user
INNER JOIN users_stats ON users_stats.id = leaderboard_%[1]s.user
%[2]s`

// LeaderboardGET gets the leaderboard.
func LeaderboardGET(md common.MethodData) common.CodeMessager {
	m := getMode(md.Query("mode"))
	w := &common.WhereClause{
		Clause: "WHERE " + md.User.OnlyUserPublic(md.HasQuery("see_everything")),
	}
	w.Where("users_stats.country = ?", md.Query("country"))
	// Admins may not want to see banned users on the leaderboard.
	// This is the default setting. In case they do, they have to activate see_everything.
	query := fmt.Sprintf(lbUserQuery, m, w.Clause+
		` ORDER BY leaderboard_`+m+`.position `+common.Paginate(md.Query("p"), md.Query("l"), 500))
	rows, err := md.DB.Query(query, w.Params...)
	if err != nil {
		md.Err(err)
		return Err500
	}
	var resp leaderboardResponse
	for rows.Next() {
		var u leaderboardUser
		err := rows.Scan(
			&u.ID, &u.Username, &u.RegisteredOn, &u.Privileges, &u.LatestActivity,

			&u.UsernameAKA, &u.Country, &u.PlayStyle, &u.FavouriteMode,

			&u.ChosenMode.RankedScore, &u.ChosenMode.TotalScore, &u.ChosenMode.PlayCount,
			&u.ChosenMode.ReplaysWatched, &u.ChosenMode.TotalHits,
			&u.ChosenMode.Accuracy, &u.ChosenMode.PP, &u.ChosenMode.GlobalLeaderboardRank,
		)
		if err != nil {
			md.Err(err)
			continue
		}
		u.ChosenMode.Level = ocl.GetLevelPrecise(int64(u.ChosenMode.TotalScore))
		resp.Users = append(resp.Users, u)
	}
	resp.Code = 200
	return resp
}
<|endoftext|>


<commit_before>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gtime_test

import (
	"github.com/gogf/gf/frame/g"
	"testing"
	"time"

	"github.com/gogf/gf/os/gtime"
	"github.com/gogf/gf/test/gtest"
)

// DATA RACE!
//time.Now()
//    /home/travis/.gimme/versions/go1.11.13.linux.amd64/src/time/time.go:1060 +0xcf
//time.sendTime()
//    /home/travis/.gimme/versions/go1.11.13.linux.amd64/src/time/sleep.go:141 +0x44
//func Test_SetTimeZone(t *testing.T) {
//	gtest.C(t, func(t *gtest.T) {
//		gtime.SetTimeZone("Asia/Shanghai")
//		t.Assert(time.Local.String(), "Asia/Shanghai")
//	})
//}

func Test_Nanosecond(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		nanos := gtime.TimestampNano()
		timeTemp := time.Unix(0, nanos)
		t.Assert(nanos, timeTemp.UnixNano())
	})
}

func Test_Microsecond(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		micros := gtime.TimestampMicro()
		timeTemp := time.Unix(0, micros*1e3)
		t.Assert(micros, timeTemp.UnixNano()/1e3)
	})
}

func Test_Millisecond(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		millis := gtime.TimestampMilli()
		timeTemp := time.Unix(0, millis*1e6)
		t.Assert(millis, timeTemp.UnixNano()/1e6)
	})
}

func Test_Second(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		s := gtime.Timestamp()
		timeTemp := time.Unix(s, 0)
		t.Assert(s, timeTemp.Unix())
	})
}

func Test_Date(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		t.Assert(gtime.Date(), time.Now().Format("2006-01-02"))
	})
}

func Test_Datetime(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		datetime := gtime.Datetime()
		timeTemp, err := gtime.StrToTime(datetime, "Y-m-d H:i:s")
		if err != nil {
			t.Error("test fail")
		}
		t.Assert(datetime, timeTemp.Time.Format("2006-01-02 15:04:05"))
	})
}

func Test_ISO8601(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		iso8601 := gtime.ISO8601()
		t.Assert(iso8601, gtime.Now().Format("c"))
	})
}

func Test_RFC822(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		rfc822 := gtime.RFC822()
		t.Assert(rfc822, gtime.Now().Format("r"))
	})
}

func Test_StrToTime(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		// Correct datetime string.
		var testDateTimes = []string{
			"2006-01-02 15:04:05",
			"2006/01/02 15:04:05",
			"2006.01.02 15:04:05.000",
			"2006.01.02 - 15:04:05",
			"2006.01.02 15:04:05 +0800 CST",
			"2006-01-02T20:05:06+05:01:01",
			"2006-01-02T14:03:04Z01:01:01",
			"2006-01-02T15:04:05Z",
			"02-jan-2006 15:04:05",
			"02/jan/2006 15:04:05",
			"02.jan.2006 15:04:05",
			"02.jan.2006:15:04:05",
		}

		for _, item := range testDateTimes {
			timeTemp, err := gtime.StrToTime(item)
			t.Assert(err, nil)
			t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")
		}

		// Correct date string,.
		var testDates = []string{
			"2006.01.02",
			"2006.01.02 00:00",
			"2006.01.02 00:00:00.000",
		}

		for _, item := range testDates {
			timeTemp, err := gtime.StrToTime(item)
			t.Assert(err, nil)
			t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05"), "2006-01-02 00:00:00")
		}

		// Correct time string.
		var testTimes = g.MapStrStr{
			"16:12:01":     "15:04:05",
			"16:12:01.789": "15:04:05.000",
		}

		for k, v := range testTimes {
			time1, err := gtime.StrToTime(k)
			t.Assert(err, nil)
			time2, err := time.ParseInLocation(v, k, time.Local)
			t.Assert(err, nil)
			t.Assert(time1.Time, time2)
		}

		// formatToStdLayout
		var testDateFormats = []string{
			"Y-m-d H:i:s",
			"\\T\\i\\m\\e Y-m-d H:i:s",
			"Y-m-d H:i:s\\",
			"Y-m-j G:i:s.u",
			"Y-m-j G:i:su",
		}

		var testDateFormatsResult = []string{
			"2007-01-02 15:04:05",
			"Time 2007-01-02 15:04:05",
			"2007-01-02 15:04:05",
			"2007-01-02 15:04:05.000",
			"2007-01-02 15:04:05.000",
		}

		for index, item := range testDateFormats {
			timeTemp, err := gtime.StrToTime(testDateFormatsResult[index], item)
			if err != nil {
				t.Error("test fail")
			}
			t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05.000"), "2007-01-02 15:04:05.000")
		}

		// 异常日期列表
		var testDatesFail = []string{
			"2006.01",
			"06..02",
		}

		for _, item := range testDatesFail {
			_, err := gtime.StrToTime(item)
			if err == nil {
				t.Error("test fail")
			}
		}

		//test err
		_, err := gtime.StrToTime("2006-01-02 15:04:05", "aabbccdd")
		if err == nil {
			t.Error("test fail")
		}
	})
}

func Test_ConvertZone(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		//现行时间
		nowUTC := time.Now().UTC()
		testZone := "America/Los_Angeles"

		//转换为洛杉矶时间
		t1, err := gtime.ConvertZone(nowUTC.Format("2006-01-02 15:04:05"), testZone, "")
		if err != nil {
			t.Error("test fail")
		}

		//使用洛杉矶时区解析上面转换后的时间
		laStr := t1.Time.Format("2006-01-02 15:04:05")
		loc, err := time.LoadLocation(testZone)
		t2, err := time.ParseInLocation("2006-01-02 15:04:05", laStr, loc)

		//判断是否与现行时间匹配
		t.Assert(t2.UTC().Unix(), nowUTC.Unix())

	})

	//test err
	gtest.C(t, func(t *gtest.T) {
		//现行时间
		nowUTC := time.Now().UTC()
		//t.Log(nowUTC.Unix())
		testZone := "errZone"

		//错误时间输入
		_, err := gtime.ConvertZone(nowUTC.Format("06..02 15:04:05"), testZone, "")
		if err == nil {
			t.Error("test fail")
		}
		//错误时区输入
		_, err = gtime.ConvertZone(nowUTC.Format("2006-01-02 15:04:05"), testZone, "")
		if err == nil {
			t.Error("test fail")
		}
		//错误时区输入
		_, err = gtime.ConvertZone(nowUTC.Format("2006-01-02 15:04:05"), testZone, testZone)
		if err == nil {
			t.Error("test fail")
		}
	})
}

func Test_ParseDuration(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("1d")
		t.Assert(err, nil)
		t.Assert(d.String(), "24h0m0s")
	})
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("1d2h3m")
		t.Assert(err, nil)
		t.Assert(d.String(), "26h3m0s")
	})
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("-1d2h3m")
		t.Assert(err, nil)
		t.Assert(d.String(), "-26h3m0s")
	})
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("3m")
		t.Assert(err, nil)
		t.Assert(d.String(), "3m0s")
	})
	// error
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("-1dd2h3m")
		t.AssertNE(err, nil)
		t.Assert(d.String(), "0s")
	})
}

func Test_ParseTimeFromContent(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		timeTemp := gtime.ParseTimeFromContent("我是中文2006-01-02 15:04:05我也是中文", "Y-m-d H:i:s")
		t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")

		timeTemp1 := gtime.ParseTimeFromContent("我是中文2006-01-02 15:04:05我也是中文")
		t.Assert(timeTemp1.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")

		timeTemp2 := gtime.ParseTimeFromContent("我是中文02.jan.2006 15:04:05我也是中文")
		t.Assert(timeTemp2.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")

		//test err
		timeTempErr := gtime.ParseTimeFromContent("我是中文", "Y-m-d H:i:s")
		if timeTempErr != nil {
			t.Error("test fail")
		}
	})
}

func Test_FuncCost(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		gtime.FuncCost(func() {

		})
	})
}
<commit_msg>fix issue of data race in unit testing case for package gtime<commit_after>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gtime_test

import (
	"github.com/gogf/gf/frame/g"
	"testing"
	"time"

	"github.com/gogf/gf/os/gtime"
	"github.com/gogf/gf/test/gtest"
)

func init() {
	gtime.SetTimeZone("Asia/Shanghai")
}
func Test_SetTimeZone(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		t.Assert(time.Local.String(), "Asia/Shanghai")
	})
}

func Test_Nanosecond(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		nanos := gtime.TimestampNano()
		timeTemp := time.Unix(0, nanos)
		t.Assert(nanos, timeTemp.UnixNano())
	})
}

func Test_Microsecond(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		micros := gtime.TimestampMicro()
		timeTemp := time.Unix(0, micros*1e3)
		t.Assert(micros, timeTemp.UnixNano()/1e3)
	})
}

func Test_Millisecond(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		millis := gtime.TimestampMilli()
		timeTemp := time.Unix(0, millis*1e6)
		t.Assert(millis, timeTemp.UnixNano()/1e6)
	})
}

func Test_Second(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		s := gtime.Timestamp()
		timeTemp := time.Unix(s, 0)
		t.Assert(s, timeTemp.Unix())
	})
}

func Test_Date(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		t.Assert(gtime.Date(), time.Now().Format("2006-01-02"))
	})
}

func Test_Datetime(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		datetime := gtime.Datetime()
		timeTemp, err := gtime.StrToTime(datetime, "Y-m-d H:i:s")
		if err != nil {
			t.Error("test fail")
		}
		t.Assert(datetime, timeTemp.Time.Format("2006-01-02 15:04:05"))
	})
}

func Test_ISO8601(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		iso8601 := gtime.ISO8601()
		t.Assert(iso8601, gtime.Now().Format("c"))
	})
}

func Test_RFC822(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		rfc822 := gtime.RFC822()
		t.Assert(rfc822, gtime.Now().Format("r"))
	})
}

func Test_StrToTime(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		// Correct datetime string.
		var testDateTimes = []string{
			"2006-01-02 15:04:05",
			"2006/01/02 15:04:05",
			"2006.01.02 15:04:05.000",
			"2006.01.02 - 15:04:05",
			"2006.01.02 15:04:05 +0800 CST",
			"2006-01-02T20:05:06+05:01:01",
			"2006-01-02T14:03:04Z01:01:01",
			"2006-01-02T15:04:05Z",
			"02-jan-2006 15:04:05",
			"02/jan/2006 15:04:05",
			"02.jan.2006 15:04:05",
			"02.jan.2006:15:04:05",
		}

		for _, item := range testDateTimes {
			timeTemp, err := gtime.StrToTime(item)
			t.Assert(err, nil)
			t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")
		}

		// Correct date string,.
		var testDates = []string{
			"2006.01.02",
			"2006.01.02 00:00",
			"2006.01.02 00:00:00.000",
		}

		for _, item := range testDates {
			timeTemp, err := gtime.StrToTime(item)
			t.Assert(err, nil)
			t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05"), "2006-01-02 00:00:00")
		}

		// Correct time string.
		var testTimes = g.MapStrStr{
			"16:12:01":     "15:04:05",
			"16:12:01.789": "15:04:05.000",
		}

		for k, v := range testTimes {
			time1, err := gtime.StrToTime(k)
			t.Assert(err, nil)
			time2, err := time.ParseInLocation(v, k, time.Local)
			t.Assert(err, nil)
			t.Assert(time1.Time, time2)
		}

		// formatToStdLayout
		var testDateFormats = []string{
			"Y-m-d H:i:s",
			"\\T\\i\\m\\e Y-m-d H:i:s",
			"Y-m-d H:i:s\\",
			"Y-m-j G:i:s.u",
			"Y-m-j G:i:su",
		}

		var testDateFormatsResult = []string{
			"2007-01-02 15:04:05",
			"Time 2007-01-02 15:04:05",
			"2007-01-02 15:04:05",
			"2007-01-02 15:04:05.000",
			"2007-01-02 15:04:05.000",
		}

		for index, item := range testDateFormats {
			timeTemp, err := gtime.StrToTime(testDateFormatsResult[index], item)
			if err != nil {
				t.Error("test fail")
			}
			t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05.000"), "2007-01-02 15:04:05.000")
		}

		// 异常日期列表
		var testDatesFail = []string{
			"2006.01",
			"06..02",
		}

		for _, item := range testDatesFail {
			_, err := gtime.StrToTime(item)
			if err == nil {
				t.Error("test fail")
			}
		}

		//test err
		_, err := gtime.StrToTime("2006-01-02 15:04:05", "aabbccdd")
		if err == nil {
			t.Error("test fail")
		}
	})
}

func Test_ConvertZone(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		//现行时间
		nowUTC := time.Now().UTC()
		testZone := "America/Los_Angeles"

		//转换为洛杉矶时间
		t1, err := gtime.ConvertZone(nowUTC.Format("2006-01-02 15:04:05"), testZone, "")
		if err != nil {
			t.Error("test fail")
		}

		//使用洛杉矶时区解析上面转换后的时间
		laStr := t1.Time.Format("2006-01-02 15:04:05")
		loc, err := time.LoadLocation(testZone)
		t2, err := time.ParseInLocation("2006-01-02 15:04:05", laStr, loc)

		//判断是否与现行时间匹配
		t.Assert(t2.UTC().Unix(), nowUTC.Unix())

	})

	//test err
	gtest.C(t, func(t *gtest.T) {
		//现行时间
		nowUTC := time.Now().UTC()
		//t.Log(nowUTC.Unix())
		testZone := "errZone"

		//错误时间输入
		_, err := gtime.ConvertZone(nowUTC.Format("06..02 15:04:05"), testZone, "")
		if err == nil {
			t.Error("test fail")
		}
		//错误时区输入
		_, err = gtime.ConvertZone(nowUTC.Format("2006-01-02 15:04:05"), testZone, "")
		if err == nil {
			t.Error("test fail")
		}
		//错误时区输入
		_, err = gtime.ConvertZone(nowUTC.Format("2006-01-02 15:04:05"), testZone, testZone)
		if err == nil {
			t.Error("test fail")
		}
	})
}

func Test_ParseDuration(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("1d")
		t.Assert(err, nil)
		t.Assert(d.String(), "24h0m0s")
	})
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("1d2h3m")
		t.Assert(err, nil)
		t.Assert(d.String(), "26h3m0s")
	})
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("-1d2h3m")
		t.Assert(err, nil)
		t.Assert(d.String(), "-26h3m0s")
	})
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("3m")
		t.Assert(err, nil)
		t.Assert(d.String(), "3m0s")
	})
	// error
	gtest.C(t, func(t *gtest.T) {
		d, err := gtime.ParseDuration("-1dd2h3m")
		t.AssertNE(err, nil)
		t.Assert(d.String(), "0s")
	})
}

func Test_ParseTimeFromContent(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		timeTemp := gtime.ParseTimeFromContent("我是中文2006-01-02 15:04:05我也是中文", "Y-m-d H:i:s")
		t.Assert(timeTemp.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")

		timeTemp1 := gtime.ParseTimeFromContent("我是中文2006-01-02 15:04:05我也是中文")
		t.Assert(timeTemp1.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")

		timeTemp2 := gtime.ParseTimeFromContent("我是中文02.jan.2006 15:04:05我也是中文")
		t.Assert(timeTemp2.Time.Format("2006-01-02 15:04:05"), "2006-01-02 15:04:05")

		//test err
		timeTempErr := gtime.ParseTimeFromContent("我是中文", "Y-m-d H:i:s")
		if timeTempErr != nil {
			t.Error("test fail")
		}
	})
}

func Test_FuncCost(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		gtime.FuncCost(func() {

		})
	})
}
<|endoftext|>


<commit_before><commit_msg>Added the IHC parameters as optimizable, for real. Made the output dir contain final results as both plot and JSON. Made the x values not used (e.g. when running BM outputs) still get the correct default values.<commit_after><|endoftext|>


<commit_before>// +build qemu
// +build network

// We only run these tests when network is activated, as the package can't run
// in parallel with QEMU engine tests. It'll also be fully covered by QEMU
// engine tests, so it's not like we strictly need to run this very often.
// If running all tests use ^go test -p 1` to ensure that multiple packages
// don't run in parallel.

package network

import (
	"net/http"
	"testing"
)

func TestNetworkCreateDestroy(t *testing.T) {
	for i := 0; i < 2; i++ {
		debug(" - Creating network pool")
		p, err := NewPool(3)
		nilOrPanic(err, "Failed to create pool")

		n1, err := p.Network()
		nilOrPanic(err, "Failed to get network")
		n2, err := p.Network()
		nilOrPanic(err, "Failed to get network")
		n3, err := p.Network()
		nilOrPanic(err, "Failed to get network")
		_, err = p.Network()
		assert(err == ErrAllNetworksInUse, "Expected ErrAllNetworksInUse")

		// Let's make a request to metaDataIP and get a 400 error
		req, err := http.NewRequest(http.MethodGet, "http://"+metaDataIP, nil)
		nilOrPanic(err, "Failed to create http request")
		res, err := http.DefaultClient.Do(req)
		nilOrPanic(err, "Failed to do http request")
		assert(res.StatusCode == http.StatusForbidden, "Expected forbidden")
		res.Body.Close()

		n1.Release()
		n1, err = p.Network()
		nilOrPanic(err, "Failed to get network")

		n1.Release()
		n2.Release()
		n3.Release()

		debug(" - Destroying network pool")
		err = p.Dispose()
		nilOrPanic(err, "Failed to dispose networks.")

		debug(" - Network pool destroyed")
	}
}
<commit_msg>Fixing more tests following???<commit_after>// +build qemu
// +build network

// We only run these tests when network is activated, as the package can't run
// in parallel with QEMU engine tests. It'll also be fully covered by QEMU
// engine tests, so it's not like we strictly need to run this very often.
// If running all tests use ^go test -p 1` to ensure that multiple packages
// don't run in parallel.

package network

import (
	"net/http"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestNetworkCreateDestroy(t *testing.T) {
	for i := 0; i < 2; i++ {
		debug(" - Creating network pool")
		p, err := NewPool(3)
		require.NoError(t, err, "Failed to create pool")

		n1, err := p.Network()
		require.NoError(t, err, "Failed to get network")
		n2, err := p.Network()
		require.NoError(t, err, "Failed to get network")
		n3, err := p.Network()
		require.NoError(t, err, "Failed to get network")
		_, err = p.Network()
		require.True(t, err == ErrAllNetworksInUse, "Expected ErrAllNetworksInUse")

		// Let's make a request to metaDataIP and get a 400 error
		req, err := http.NewRequest(http.MethodGet, "http://"+metaDataIP, nil)
		require.NoError(t, err, "Failed to create http request")
		res, err := http.DefaultClient.Do(req)
		require.NoError(t, err, "Failed to do http request")
		require.True(t, res.StatusCode == http.StatusForbidden, "Expected forbidden")
		res.Body.Close()

		n1.Release()
		n1, err = p.Network()
		require.NoError(t, err, "Failed to get network")

		n1.Release()
		n2.Release()
		n3.Release()

		debug(" - Destroying network pool")
		err = p.Dispose()
		require.NoError(t, err, "Failed to dispose networks.")

		debug(" - Network pool destroyed")
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"os"
	"runtime"
	"testing"
)

var listenMulticastUDPTests = []struct {
	net   string
	gaddr *UDPAddr
	flags Flags
	ipv6  bool
}{
	// cf. RFC 4727: Experimental Values in IPv4, IPv6, ICMPv4, ICMPv6, UDP, and TCP Headers
	{"udp", &UDPAddr{IPv4(224, 0, 0, 254), 12345}, FlagUp | FlagLoopback, false},
	{"udp4", &UDPAddr{IPv4(224, 0, 0, 254), 12345}, FlagUp | FlagLoopback, false},
	{"udp", &UDPAddr{ParseIP("ff0e::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff01::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff02::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff04::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff05::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff08::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff0e::114"), 12345}, FlagUp | FlagLoopback, true},
}

func TestListenMulticastUDP(t *testing.T) {
	switch runtime.GOOS {
	case "netbsd", "openbsd", "plan9", "windows":
		return
	case "linux":
		if runtime.GOARCH == "arm" {
			return
		}
	}

	for _, tt := range listenMulticastUDPTests {
		if tt.ipv6 && (!supportsIPv6 || os.Getuid() != 0) {
			continue
		}
		ift, err := Interfaces()
		if err != nil {
			t.Fatalf("Interfaces failed: %v", err)
		}
		var ifi *Interface
		for _, x := range ift {
			if x.Flags&tt.flags == tt.flags {
				ifi = &x
				break
			}
		}
		if ifi == nil {
			t.Logf("an appropriate multicast interface not found")
			return
		}
		c, err := ListenMulticastUDP(tt.net, ifi, tt.gaddr)
		if err != nil {
			t.Fatalf("ListenMulticastUDP failed: %v", err)
		}
		defer c.Close() // test to listen concurrently across multiple listeners
		if !tt.ipv6 {
			testIPv4MulticastSocketOptions(t, c.fd, ifi)
		} else {
			testIPv6MulticastSocketOptions(t, c.fd, ifi)
		}
		ifmat, err := ifi.MulticastAddrs()
		if err != nil {
			t.Fatalf("MulticastAddrs failed: %v", err)
		}
		var found bool
		for _, ifma := range ifmat {
			if ifma.(*IPAddr).IP.Equal(tt.gaddr.IP) {
				found = true
				break
			}
		}
		if !found {
			t.Fatalf("%q not found in RIB", tt.gaddr.String())
		}
	}
}

func TestSimpleListenMulticastUDP(t *testing.T) {
	switch runtime.GOOS {
	case "plan9":
		return
	}

	for _, tt := range listenMulticastUDPTests {
		if tt.ipv6 {
			continue
		}
		tt.flags = FlagUp | FlagMulticast
		ift, err := Interfaces()
		if err != nil {
			t.Fatalf("Interfaces failed: %v", err)
		}
		var ifi *Interface
		for _, x := range ift {
			if x.Flags&tt.flags == tt.flags {
				ifi = &x
				break
			}
		}
		if ifi == nil {
			t.Logf("an appropriate multicast interface not found")
			return
		}
		c, err := ListenMulticastUDP(tt.net, ifi, tt.gaddr)
		if err != nil {
			t.Fatalf("ListenMulticastUDP failed: %v", err)
		}
		c.Close()
	}
}

func testIPv4MulticastSocketOptions(t *testing.T, fd *netFD, ifi *Interface) {
	ifmc, err := ipv4MulticastInterface(fd)
	if err != nil {
		t.Fatalf("ipv4MulticastInterface failed: %v", err)
	}
	t.Logf("IPv4 multicast interface: %v", ifmc)
	err = setIPv4MulticastInterface(fd, ifi)
	if err != nil {
		t.Fatalf("setIPv4MulticastInterface failed: %v", err)
	}

	ttl, err := ipv4MulticastTTL(fd)
	if err != nil {
		t.Fatalf("ipv4MulticastTTL failed: %v", err)
	}
	t.Logf("IPv4 multicast TTL: %v", ttl)
	err = setIPv4MulticastTTL(fd, 1)
	if err != nil {
		t.Fatalf("setIPv4MulticastTTL failed: %v", err)
	}

	loop, err := ipv4MulticastLoopback(fd)
	if err != nil {
		t.Fatalf("ipv4MulticastLoopback failed: %v", err)
	}
	t.Logf("IPv4 multicast loopback: %v", loop)
	err = setIPv4MulticastLoopback(fd, false)
	if err != nil {
		t.Fatalf("setIPv4MulticastLoopback failed: %v", err)
	}
}

func testIPv6MulticastSocketOptions(t *testing.T, fd *netFD, ifi *Interface) {
	ifmc, err := ipv6MulticastInterface(fd)
	if err != nil {
		t.Fatalf("ipv6MulticastInterface failed: %v", err)
	}
	t.Logf("IPv6 multicast interface: %v", ifmc)
	err = setIPv6MulticastInterface(fd, ifi)
	if err != nil {
		t.Fatalf("setIPv6MulticastInterface failed: %v", err)
	}

	hoplim, err := ipv6MulticastHopLimit(fd)
	if err != nil {
		t.Fatalf("ipv6MulticastHopLimit failed: %v", err)
	}
	t.Logf("IPv6 multicast hop limit: %v", hoplim)
	err = setIPv6MulticastHopLimit(fd, 1)
	if err != nil {
		t.Fatalf("setIPv6MulticastHopLimit failed: %v", err)
	}

	loop, err := ipv6MulticastLoopback(fd)
	if err != nil {
		t.Fatalf("ipv6MulticastLoopback failed: %v", err)
	}
	t.Logf("IPv6 multicast loopback: %v", loop)
	err = setIPv6MulticastLoopback(fd, false)
	if err != nil {
		t.Fatalf("setIPv6MulticastLoopback failed: %v", err)
	}
}
<commit_msg>net: disable multicast test on Alpha GNU/Linux<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"os"
	"runtime"
	"testing"
)

var listenMulticastUDPTests = []struct {
	net   string
	gaddr *UDPAddr
	flags Flags
	ipv6  bool
}{
	// cf. RFC 4727: Experimental Values in IPv4, IPv6, ICMPv4, ICMPv6, UDP, and TCP Headers
	{"udp", &UDPAddr{IPv4(224, 0, 0, 254), 12345}, FlagUp | FlagLoopback, false},
	{"udp4", &UDPAddr{IPv4(224, 0, 0, 254), 12345}, FlagUp | FlagLoopback, false},
	{"udp", &UDPAddr{ParseIP("ff0e::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff01::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff02::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff04::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff05::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff08::114"), 12345}, FlagUp | FlagLoopback, true},
	{"udp6", &UDPAddr{ParseIP("ff0e::114"), 12345}, FlagUp | FlagLoopback, true},
}

func TestListenMulticastUDP(t *testing.T) {
	switch runtime.GOOS {
	case "netbsd", "openbsd", "plan9", "windows":
		return
	case "linux":
		if runtime.GOARCH == "arm" || runtime.GOARCH == "alpha" {
			return
		}
	}

	for _, tt := range listenMulticastUDPTests {
		if tt.ipv6 && (!supportsIPv6 || os.Getuid() != 0) {
			continue
		}
		ift, err := Interfaces()
		if err != nil {
			t.Fatalf("Interfaces failed: %v", err)
		}
		var ifi *Interface
		for _, x := range ift {
			if x.Flags&tt.flags == tt.flags {
				ifi = &x
				break
			}
		}
		if ifi == nil {
			t.Logf("an appropriate multicast interface not found")
			return
		}
		c, err := ListenMulticastUDP(tt.net, ifi, tt.gaddr)
		if err != nil {
			t.Fatalf("ListenMulticastUDP failed: %v", err)
		}
		defer c.Close() // test to listen concurrently across multiple listeners
		if !tt.ipv6 {
			testIPv4MulticastSocketOptions(t, c.fd, ifi)
		} else {
			testIPv6MulticastSocketOptions(t, c.fd, ifi)
		}
		ifmat, err := ifi.MulticastAddrs()
		if err != nil {
			t.Fatalf("MulticastAddrs failed: %v", err)
		}
		var found bool
		for _, ifma := range ifmat {
			if ifma.(*IPAddr).IP.Equal(tt.gaddr.IP) {
				found = true
				break
			}
		}
		if !found {
			t.Fatalf("%q not found in RIB", tt.gaddr.String())
		}
	}
}

func TestSimpleListenMulticastUDP(t *testing.T) {
	switch runtime.GOOS {
	case "plan9":
		return
	}

	for _, tt := range listenMulticastUDPTests {
		if tt.ipv6 {
			continue
		}
		tt.flags = FlagUp | FlagMulticast
		ift, err := Interfaces()
		if err != nil {
			t.Fatalf("Interfaces failed: %v", err)
		}
		var ifi *Interface
		for _, x := range ift {
			if x.Flags&tt.flags == tt.flags {
				ifi = &x
				break
			}
		}
		if ifi == nil {
			t.Logf("an appropriate multicast interface not found")
			return
		}
		c, err := ListenMulticastUDP(tt.net, ifi, tt.gaddr)
		if err != nil {
			t.Fatalf("ListenMulticastUDP failed: %v", err)
		}
		c.Close()
	}
}

func testIPv4MulticastSocketOptions(t *testing.T, fd *netFD, ifi *Interface) {
	ifmc, err := ipv4MulticastInterface(fd)
	if err != nil {
		t.Fatalf("ipv4MulticastInterface failed: %v", err)
	}
	t.Logf("IPv4 multicast interface: %v", ifmc)
	err = setIPv4MulticastInterface(fd, ifi)
	if err != nil {
		t.Fatalf("setIPv4MulticastInterface failed: %v", err)
	}

	ttl, err := ipv4MulticastTTL(fd)
	if err != nil {
		t.Fatalf("ipv4MulticastTTL failed: %v", err)
	}
	t.Logf("IPv4 multicast TTL: %v", ttl)
	err = setIPv4MulticastTTL(fd, 1)
	if err != nil {
		t.Fatalf("setIPv4MulticastTTL failed: %v", err)
	}

	loop, err := ipv4MulticastLoopback(fd)
	if err != nil {
		t.Fatalf("ipv4MulticastLoopback failed: %v", err)
	}
	t.Logf("IPv4 multicast loopback: %v", loop)
	err = setIPv4MulticastLoopback(fd, false)
	if err != nil {
		t.Fatalf("setIPv4MulticastLoopback failed: %v", err)
	}
}

func testIPv6MulticastSocketOptions(t *testing.T, fd *netFD, ifi *Interface) {
	ifmc, err := ipv6MulticastInterface(fd)
	if err != nil {
		t.Fatalf("ipv6MulticastInterface failed: %v", err)
	}
	t.Logf("IPv6 multicast interface: %v", ifmc)
	err = setIPv6MulticastInterface(fd, ifi)
	if err != nil {
		t.Fatalf("setIPv6MulticastInterface failed: %v", err)
	}

	hoplim, err := ipv6MulticastHopLimit(fd)
	if err != nil {
		t.Fatalf("ipv6MulticastHopLimit failed: %v", err)
	}
	t.Logf("IPv6 multicast hop limit: %v", hoplim)
	err = setIPv6MulticastHopLimit(fd, 1)
	if err != nil {
		t.Fatalf("setIPv6MulticastHopLimit failed: %v", err)
	}

	loop, err := ipv6MulticastLoopback(fd)
	if err != nil {
		t.Fatalf("ipv6MulticastLoopback failed: %v", err)
	}
	t.Logf("IPv6 multicast loopback: %v", loop)
	err = setIPv6MulticastLoopback(fd, false)
	if err != nil {
		t.Fatalf("setIPv6MulticastLoopback failed: %v", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"os/exec"
	"strconv"
	"strings"

	"github.com/NetSys/quilt/api"
	"github.com/NetSys/quilt/api/client/getter"
	"github.com/NetSys/quilt/db"

	log "github.com/Sirupsen/logrus"
)

func main() {
	clientGetter := getter.New()

	clnt, err := clientGetter.Client(api.DefaultSocket)
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't get quiltctl client")
	}
	defer clnt.Close()

	leader, err := clientGetter.LeaderClient(clnt)
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't get leader client")
	}

	containers, err := leader.QueryContainers()
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't query containers")
	}

	machines, err := clnt.QueryMachines()
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't query machines")
	}

	if logContainers(containers) && httpGetTest(machines, containers) {
		log.Info("PASSED")
	} else {
		log.Info("FAILED")
	}
}

func logContainers(containers []db.Container) bool {
	var failed bool
	for _, c := range containers {
		cmd := exec.Command("quilt", "logs", strconv.Itoa(c.StitchID))
		out, err := cmd.CombinedOutput()
		if err != nil {
			log.WithError(err).Error("Failed to log: %s", c)
			failed = true
		}
		log.Infof("Container: %s\n%s\n\n", c, string(out))
	}

	return !failed
}

func httpGetTest(machines []db.Machine, containers []db.Container) bool {
	log.Info("HTTP Get Test")

	minionIPMap := map[string]string{}
	for _, m := range machines {
		minionIPMap[m.PrivateIP] = m.PublicIP
	}

	var publicIPs []string
	for _, c := range containers {
		if strings.Contains(c.Image, "haproxy") {
			ip, ok := minionIPMap[c.Minion]
			if !ok {
				log.WithField("container", c).Fatal(
					"FAILED, HAProxy with no public IP")
			}
			publicIPs = append(publicIPs, ip)
		}
	}

	log.Info("Public IPs: ", publicIPs)
	if len(publicIPs) == 0 {
		log.Fatal("FAILED, Found no public IPs")
	}

	var failed bool
	for i := 0; i < 25; i++ {
		for _, ip := range publicIPs {
			resp, err := http.Get("http://" + ip)
			if err != nil {
				log.WithError(err).Error("HTTP Error")
				failed = true
			}

			if resp.StatusCode == 200 {
				log.Info(resp)
			} else {
				log.Error(resp)
				failed = true
			}
		}
	}

	return !failed
}
<commit_msg>tester: Fix a panic in the mean test<commit_after>package main

import (
	"net/http"
	"os/exec"
	"strconv"
	"strings"

	"github.com/NetSys/quilt/api"
	"github.com/NetSys/quilt/api/client/getter"
	"github.com/NetSys/quilt/db"

	log "github.com/Sirupsen/logrus"
)

func main() {
	clientGetter := getter.New()

	clnt, err := clientGetter.Client(api.DefaultSocket)
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't get quiltctl client")
	}
	defer clnt.Close()

	leader, err := clientGetter.LeaderClient(clnt)
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't get leader client")
	}

	containers, err := leader.QueryContainers()
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't query containers")
	}

	machines, err := clnt.QueryMachines()
	if err != nil {
		log.WithError(err).Fatal("FAILED, couldn't query machines")
	}

	if logContainers(containers) && httpGetTest(machines, containers) {
		log.Info("PASSED")
	} else {
		log.Info("FAILED")
	}
}

func logContainers(containers []db.Container) bool {
	var failed bool
	for _, c := range containers {
		cmd := exec.Command("quilt", "logs", strconv.Itoa(c.StitchID))
		out, err := cmd.CombinedOutput()
		if err != nil {
			log.WithError(err).Error("Failed to log: %s", c)
			failed = true
		}
		log.Infof("Container: %s\n%s\n\n", c, string(out))
	}

	return !failed
}

func httpGetTest(machines []db.Machine, containers []db.Container) bool {
	log.Info("HTTP Get Test")

	minionIPMap := map[string]string{}
	for _, m := range machines {
		minionIPMap[m.PrivateIP] = m.PublicIP
	}

	var publicIPs []string
	for _, c := range containers {
		if strings.Contains(c.Image, "haproxy") {
			ip, ok := minionIPMap[c.Minion]
			if !ok {
				log.WithField("container", c).Fatal(
					"FAILED, HAProxy with no public IP")
			}
			publicIPs = append(publicIPs, ip)
		}
	}

	log.Info("Public IPs: ", publicIPs)
	if len(publicIPs) == 0 {
		log.Fatal("FAILED, Found no public IPs")
	}

	var failed bool
	for i := 0; i < 25; i++ {
		for _, ip := range publicIPs {
			resp, err := http.Get("http://" + ip)
			if err != nil {
				log.WithError(err).Error("HTTP Error")
				failed = true
				continue
			}

			if resp.StatusCode == 200 {
				log.Info(resp)
			} else {
				log.Error(resp)
				failed = true
			}
		}
	}

	return !failed
}
<|endoftext|>


<commit_before>// +build linux

package sysstat

import (
	"bufio"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

var (
	firstTimestamp time.Time
	execCommand    = exec.Command // execCommand is used to mock commands in tests.
	dfltActivities = []string{"DISK"}
)

const parseInterval = 1 // parseInterval is the interval (in seconds) where the parsing of the binary file takes place.

type Sysstat struct {
	// Sadc represents the path to the sadc collector utility.
	Sadc string `toml:"sadc_path"`

	// Sadf represents the path to the sadf cmd.
	Sadf string `toml:"sadf_path"`

	// Activities is a list of activities that are passed as argument to the
	// collector utility (e.g: DISK, SNMP etc...)
	// The more activities that are added, the more data is collected.
	Activities []string

	// Options is a map of options.
	//
	// The key represents the actual option that the Sadf command is called with and
	// the value represents the description for that option.
	//
	// For example, if you have the following options map:
	//    map[string]string{"-C": "cpu", "-d": "disk"}
	// The Sadf command is run with the options -C and -d to extract cpu and
	// disk metrics from the collected binary file.
	//
	// If Group is false (see below), each metric will be prefixed with the corresponding description
	// and represents itself a measurement.
	//
	// If Group is true, metrics are grouped to a single measurement with the corresponding description as name.
	Options map[string]string

	// Group determines if metrics are grouped or not.
	Group bool

	// DeviceTags adds the possibility to add additional tags for devices.
	DeviceTags map[string][]map[string]string `toml:"device_tags"`
	tmpFile    string
	interval   int
}

func (*Sysstat) Description() string {
	return "Sysstat metrics collector"
}

var sampleConfig = `
  ## Path to the sadc command.
  #
  ## On Debian and Arch Linux the default path is /usr/lib/sa/sadc whereas
  ## on RHEL and CentOS the default path is /usr/lib64/sa/sadc
  sadc_path = "/usr/lib/sa/sadc" # required
  #
  #
  ## Path to the sadf command, if it is not in PATH
  # sadf_path = "/usr/bin/sadf"
  #
  #
  ## Activities is a list of activities, that are passed as argument to the
  ## sadc collector utility (e.g: DISK, SNMP etc...)
  ## The more activities that are added, the more data is collected.
  # activities = ["DISK"]
  #
  #
  ## Group metrics to measurements.
  ##
  ## If group is false each metric will be prefixed with a description
  ## and represents itself a measurement.
  ##
  ## If Group is true, corresponding metrics are grouped to a single measurement.
  # group = false
  #
  #
  ## Options for the sadf command. The values on the left represent the sadf options and
  ## the values on the right their description (wich are used for grouping and prefixing metrics).
  [inputs.sysstat.options]
	-C = "cpu"
	-B = "paging"
	-b = "io"
	-d = "disk"             # requires DISK activity
	-H = "hugepages"
	"-n ALL" = "network"
	"-P ALL" = "per_cpu"
	-q = "queue"
	-R = "mem"
	"-r ALL" = "mem_util"
	-S = "swap_util"
	-u = "cpu_util"
	-v = "inode"
	-W = "swap"
	-w = "task"
  #	"-I ALL" = "interrupts" # requires INT activity
  #
  #
  ## Device tags can be used to add additional tags for devices. For example the configuration below
  ## adds a tag vg with value rootvg for all metrics with sda devices.
  # [[inputs.sysstat.device_tags.sda]]
  #  vg = "rootvg"
`

func (*Sysstat) SampleConfig() string {
	return sampleConfig
}

func (s *Sysstat) Gather(acc telegraf.Accumulator) error {
	if s.interval == 0 {
		if firstTimestamp.IsZero() {
			firstTimestamp = time.Now()
		} else {
			s.interval = int(time.Since(firstTimestamp).Seconds())
		}
	}
	ts := time.Now().Add(time.Duration(s.interval) * time.Second)
	if err := s.collect(); err != nil {
		return err
	}
	var wg sync.WaitGroup
	errorChannel := make(chan error, len(s.Options)*2)
	for option := range s.Options {
		wg.Add(1)
		go func(acc telegraf.Accumulator, option string) {
			defer wg.Done()
			if err := s.parse(acc, option, ts); err != nil {
				errorChannel <- err
			}
		}(acc, option)
	}
	wg.Wait()
	close(errorChannel)

	errorStrings := []string{}
	for err := range errorChannel {
		errorStrings = append(errorStrings, err.Error())
	}

	if _, err := os.Stat(s.tmpFile); err == nil {
		if err := os.Remove(s.tmpFile); err != nil {
			errorStrings = append(errorStrings, err.Error())
		}
	}

	if len(errorStrings) == 0 {
		return nil
	}
	return errors.New(strings.Join(errorStrings, "\n"))
}

// collect collects sysstat data with the collector utility sadc. It runs the following command:
//     Sadc -S <Activity1> -S <Activity2> ... <collectInterval> 2 tmpFile
// The above command collects system metrics during <collectInterval> and saves it in binary form to tmpFile.
func (s *Sysstat) collect() error {
	if len(s.Activities) == 0 {
		s.Activities = dfltActivities
	}
	if len(s.Sadf) == 0 {
		sadf, err := exec.LookPath("sadf")
		if err != nil {
			return errors.New("sadf not in $PATH, configure path to sadf")
		}
		s.Sadf = sadf
	}
	options := []string{}
	for _, act := range s.Activities {
		options = append(options, "-S", act)
	}
	s.tmpFile = path.Join("/tmp", fmt.Sprintf("sysstat-%d", time.Now().Unix()))
	collectInterval := s.interval - parseInterval // collectInterval has to be smaller than the telegraf data collection interval

	if collectInterval < 0 { // If true, interval is not defined yet and Gather is run for the first time.
		collectInterval = 1 // In that case we only collect for 1 second.
	}

	options = append(options, strconv.Itoa(collectInterval), "2", s.tmpFile)
	cmd := execCommand(s.Sadc, options...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to run command %s: %s", strings.Join(cmd.Args, " "), string(out))
	}
	return nil
}

// parse runs Sadf on the previously saved tmpFile:
//    Sadf -p -- -p <option> tmpFile
// and parses the output to add it to the telegraf.Accumulator acc.
func (s *Sysstat) parse(acc telegraf.Accumulator, option string, ts time.Time) error {
	cmd := execCommand(s.Sadf, s.sadfOptions(option)...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("running command '%s' failed: %s", strings.Join(cmd.Args, " "), err)
	}

	r := bufio.NewReader(stdout)
	csv := csv.NewReader(r)
	csv.Comma = '\t'
	csv.FieldsPerRecord = 6
	var measurement string
	// groupData to accumulate data when Group=true
	type groupData struct {
		tags   map[string]string
		fields map[string]interface{}
	}
	m := make(map[string]groupData)
	for {
		record, err := csv.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		device := record[3]
		value, err := strconv.ParseFloat(record[5], 64)
		if err != nil {
			return err
		}

		tags := map[string]string{}
		if device != "-" {
			tags["device"] = device
			if addTags, ok := s.DeviceTags[device]; ok {
				for _, tag := range addTags {
					for k, v := range tag {
						tags[k] = v
					}
				}

			}
		}

		if s.Group {
			measurement = s.Options[option]
			if _, ok := m[device]; !ok {
				m[device] = groupData{
					fields: make(map[string]interface{}),
					tags:   make(map[string]string),
				}
			}
			g, _ := m[device]
			if len(g.tags) == 0 {
				for k, v := range tags {
					g.tags[k] = v
				}
			}
			g.fields[escape(record[4])] = value
		} else {
			measurement = s.Options[option] + "_" + escape(record[4])
			fields := map[string]interface{}{
				"value": value,
			}
			acc.AddFields(measurement, fields, tags, ts)
		}

	}
	if s.Group {
		for _, v := range m {
			acc.AddFields(measurement, v.fields, v.tags, ts)
		}
	}
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("command %s failed with %s", strings.Join(cmd.Args, " "), err)
	}
	return nil
}

// sadfOptions creates the correct options for the sadf utility.
func (s *Sysstat) sadfOptions(activityOption string) []string {
	options := []string{
		"-p",
		"--",
		"-p",
	}

	opts := strings.Split(activityOption, " ")
	options = append(options, opts...)
	options = append(options, s.tmpFile)

	return options
}

// escape removes % and / chars in field names
func escape(dirty string) string {
	var fieldEscaper = strings.NewReplacer(
		`%`, "pct_",
		`/`, "_per_",
	)
	return fieldEscaper.Replace(dirty)
}

func init() {
	inputs.Add("sysstat", func() telegraf.Input {
		return &Sysstat{}
	})
}
<commit_msg>change group=true by default<commit_after>// +build linux

package sysstat

import (
	"bufio"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

var (
	firstTimestamp time.Time
	execCommand    = exec.Command // execCommand is used to mock commands in tests.
	dfltActivities = []string{"DISK"}
)

const parseInterval = 1 // parseInterval is the interval (in seconds) where the parsing of the binary file takes place.

type Sysstat struct {
	// Sadc represents the path to the sadc collector utility.
	Sadc string `toml:"sadc_path"`

	// Sadf represents the path to the sadf cmd.
	Sadf string `toml:"sadf_path"`

	// Activities is a list of activities that are passed as argument to the
	// collector utility (e.g: DISK, SNMP etc...)
	// The more activities that are added, the more data is collected.
	Activities []string

	// Options is a map of options.
	//
	// The key represents the actual option that the Sadf command is called with and
	// the value represents the description for that option.
	//
	// For example, if you have the following options map:
	//    map[string]string{"-C": "cpu", "-d": "disk"}
	// The Sadf command is run with the options -C and -d to extract cpu and
	// disk metrics from the collected binary file.
	//
	// If Group is false (see below), each metric will be prefixed with the corresponding description
	// and represents itself a measurement.
	//
	// If Group is true, metrics are grouped to a single measurement with the corresponding description as name.
	Options map[string]string

	// Group determines if metrics are grouped or not.
	Group bool

	// DeviceTags adds the possibility to add additional tags for devices.
	DeviceTags map[string][]map[string]string `toml:"device_tags"`
	tmpFile    string
	interval   int
}

func (*Sysstat) Description() string {
	return "Sysstat metrics collector"
}

var sampleConfig = `
  ## Path to the sadc command.
  #
  ## On Debian and Arch Linux the default path is /usr/lib/sa/sadc whereas
  ## on RHEL and CentOS the default path is /usr/lib64/sa/sadc
  sadc_path = "/usr/lib/sa/sadc" # required
  #
  #
  ## Path to the sadf command, if it is not in PATH
  # sadf_path = "/usr/bin/sadf"
  #
  #
  ## Activities is a list of activities, that are passed as argument to the
  ## sadc collector utility (e.g: DISK, SNMP etc...)
  ## The more activities that are added, the more data is collected.
  # activities = ["DISK"]
  #
  #
  ## Group metrics to measurements.
  ##
  ## If group is false each metric will be prefixed with a description
  ## and represents itself a measurement.
  ##
  ## If Group is true, corresponding metrics are grouped to a single measurement.
  # group = false
  #
  #
  ## Options for the sadf command. The values on the left represent the sadf options and
  ## the values on the right their description (wich are used for grouping and prefixing metrics).
  [inputs.sysstat.options]
	-C = "cpu"
	-B = "paging"
	-b = "io"
	-d = "disk"             # requires DISK activity
	-H = "hugepages"
	"-n ALL" = "network"
	"-P ALL" = "per_cpu"
	-q = "queue"
	-R = "mem"
	"-r ALL" = "mem_util"
	-S = "swap_util"
	-u = "cpu_util"
	-v = "inode"
	-W = "swap"
	-w = "task"
  #	"-I ALL" = "interrupts" # requires INT activity
  #
  #
  ## Device tags can be used to add additional tags for devices. For example the configuration below
  ## adds a tag vg with value rootvg for all metrics with sda devices.
  # [[inputs.sysstat.device_tags.sda]]
  #  vg = "rootvg"
`

func (*Sysstat) SampleConfig() string {
	return sampleConfig
}

func (s *Sysstat) Gather(acc telegraf.Accumulator) error {
	if s.interval == 0 {
		if firstTimestamp.IsZero() {
			firstTimestamp = time.Now()
		} else {
			s.interval = int(time.Since(firstTimestamp).Seconds())
		}
	}
	ts := time.Now().Add(time.Duration(s.interval) * time.Second)
	if err := s.collect(); err != nil {
		return err
	}
	var wg sync.WaitGroup
	errorChannel := make(chan error, len(s.Options)*2)
	for option := range s.Options {
		wg.Add(1)
		go func(acc telegraf.Accumulator, option string) {
			defer wg.Done()
			if err := s.parse(acc, option, ts); err != nil {
				errorChannel <- err
			}
		}(acc, option)
	}
	wg.Wait()
	close(errorChannel)

	errorStrings := []string{}
	for err := range errorChannel {
		errorStrings = append(errorStrings, err.Error())
	}

	if _, err := os.Stat(s.tmpFile); err == nil {
		if err := os.Remove(s.tmpFile); err != nil {
			errorStrings = append(errorStrings, err.Error())
		}
	}

	if len(errorStrings) == 0 {
		return nil
	}
	return errors.New(strings.Join(errorStrings, "\n"))
}

// collect collects sysstat data with the collector utility sadc. It runs the following command:
//     Sadc -S <Activity1> -S <Activity2> ... <collectInterval> 2 tmpFile
// The above command collects system metrics during <collectInterval> and saves it in binary form to tmpFile.
func (s *Sysstat) collect() error {
	if len(s.Activities) == 0 {
		s.Activities = dfltActivities
	}
	if len(s.Sadf) == 0 {
		sadf, err := exec.LookPath("sadf")
		if err != nil {
			return errors.New("sadf not in $PATH, configure path to sadf")
		}
		s.Sadf = sadf
	}
	options := []string{}
	for _, act := range s.Activities {
		options = append(options, "-S", act)
	}
	s.tmpFile = path.Join("/tmp", fmt.Sprintf("sysstat-%d", time.Now().Unix()))
	collectInterval := s.interval - parseInterval // collectInterval has to be smaller than the telegraf data collection interval

	if collectInterval < 0 { // If true, interval is not defined yet and Gather is run for the first time.
		collectInterval = 1 // In that case we only collect for 1 second.
	}

	options = append(options, strconv.Itoa(collectInterval), "2", s.tmpFile)
	cmd := execCommand(s.Sadc, options...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to run command %s: %s", strings.Join(cmd.Args, " "), string(out))
	}
	return nil
}

// parse runs Sadf on the previously saved tmpFile:
//    Sadf -p -- -p <option> tmpFile
// and parses the output to add it to the telegraf.Accumulator acc.
func (s *Sysstat) parse(acc telegraf.Accumulator, option string, ts time.Time) error {
	cmd := execCommand(s.Sadf, s.sadfOptions(option)...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("running command '%s' failed: %s", strings.Join(cmd.Args, " "), err)
	}

	r := bufio.NewReader(stdout)
	csv := csv.NewReader(r)
	csv.Comma = '\t'
	csv.FieldsPerRecord = 6
	var measurement string
	// groupData to accumulate data when Group=true
	type groupData struct {
		tags   map[string]string
		fields map[string]interface{}
	}
	m := make(map[string]groupData)
	for {
		record, err := csv.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		device := record[3]
		value, err := strconv.ParseFloat(record[5], 64)
		if err != nil {
			return err
		}

		tags := map[string]string{}
		if device != "-" {
			tags["device"] = device
			if addTags, ok := s.DeviceTags[device]; ok {
				for _, tag := range addTags {
					for k, v := range tag {
						tags[k] = v
					}
				}

			}
		}

		if s.Group {
			measurement = s.Options[option]
			if _, ok := m[device]; !ok {
				m[device] = groupData{
					fields: make(map[string]interface{}),
					tags:   make(map[string]string),
				}
			}
			g, _ := m[device]
			if len(g.tags) == 0 {
				for k, v := range tags {
					g.tags[k] = v
				}
			}
			g.fields[escape(record[4])] = value
		} else {
			measurement = s.Options[option] + "_" + escape(record[4])
			fields := map[string]interface{}{
				"value": value,
			}
			acc.AddFields(measurement, fields, tags, ts)
		}

	}
	if s.Group {
		for _, v := range m {
			acc.AddFields(measurement, v.fields, v.tags, ts)
		}
	}
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("command %s failed with %s", strings.Join(cmd.Args, " "), err)
	}
	return nil
}

// sadfOptions creates the correct options for the sadf utility.
func (s *Sysstat) sadfOptions(activityOption string) []string {
	options := []string{
		"-p",
		"--",
		"-p",
	}

	opts := strings.Split(activityOption, " ")
	options = append(options, opts...)
	options = append(options, s.tmpFile)

	return options
}

// escape removes % and / chars in field names
func escape(dirty string) string {
	var fieldEscaper = strings.NewReplacer(
		`%`, "pct_",
		`/`, "_per_",
	)
	return fieldEscaper.Replace(dirty)
}

func init() {
	inputs.Add("sysstat", func() telegraf.Input {
		return &Sysstat{
			Group: true,
		}
	})
}
<|endoftext|>


<commit_before>// package blockstore implements a thin wrapper over a datastore, giving a
// clean interface for Getting and Putting block objects.
package blockstore

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"

	blocks "github.com/ipfs/go-ipfs/blocks"
	dshelp "github.com/ipfs/go-ipfs/thirdparty/ds-help"

	ds "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore"
	dsns "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore/namespace"
	dsq "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore/query"
	logging "gx/ipfs/QmSpJByNKFX1sCsHBEp3R73FL4NF6FnQTEGyNAXHm2GS52/go-log"
	cid "gx/ipfs/QmcTcsTvfaeEBRFo1TkFgT8sRmgi1n1LTZpecfVP8fzpGD/go-cid"
)

var log = logging.Logger("blockstore")

// BlockPrefix namespaces blockstore datastores
var BlockPrefix = ds.NewKey("blocks")

var ValueTypeMismatch = errors.New("the retrieved value is not a Block")
var ErrHashMismatch = errors.New("block in storage has different hash than requested")

var ErrNotFound = errors.New("blockstore: block not found")

// Blockstore wraps a Datastore
type Blockstore interface {
	DeleteBlock(*cid.Cid) error
	Has(*cid.Cid) (bool, error)
	Get(*cid.Cid) (blocks.Block, error)
	Put(blocks.Block) error
	PutMany([]blocks.Block) error

	AllKeysChan(ctx context.Context) (<-chan *cid.Cid, error)
}

type GCLocker interface {
	// GCLock locks the blockstore for garbage collection. No operations
	// that expect to finish with a pin should ocurr simultaneously.
	// Reading during GC is safe, and requires no lock.
	GCLock() Unlocker

	// PinLock locks the blockstore for sequences of puts expected to finish
	// with a pin (before GC). Multiple put->pin sequences can write through
	// at the same time, but no GC should not happen simulatenously.
	// Reading during Pinning is safe, and requires no lock.
	PinLock() Unlocker

	// GcRequested returns true if GCLock has been called and is waiting to
	// take the lock
	GCRequested() bool
}

type GCBlockstore interface {
	Blockstore
	GCLocker
}

func NewGCBlockstore(bs Blockstore, gcl GCLocker) GCBlockstore {
	return gcBlockstore{bs, gcl}
}

type gcBlockstore struct {
	Blockstore
	GCLocker
}

func NewBlockstore(d ds.Batching) *blockstore {
	var dsb ds.Batching
	dd := dsns.Wrap(d, BlockPrefix)
	dsb = dd
	return &blockstore{
		datastore: dsb,
	}
}

type blockstore struct {
	datastore ds.Batching

	lk      sync.RWMutex
	gcreq   int32
	gcreqlk sync.Mutex

	rehash bool
}

func (bs *blockstore) HashOnRead(enabled bool) {
	bs.rehash = enabled
}

func (bs *blockstore) Get(k *cid.Cid) (blocks.Block, error) {
	if k == nil {
		log.Error("nil cid in blockstore")
		return nil, ErrNotFound
	}

	maybeData, err := bs.datastore.Get(dshelp.CidToDsKey(k))
	if err == ds.ErrNotFound {
		return nil, ErrNotFound
	}
	if err != nil {
		return nil, err
	}
	bdata, ok := maybeData.([]byte)
	if !ok {
		return nil, ValueTypeMismatch
	}

	if bs.rehash {
		rbcid, err := k.Prefix().Sum(bdata)
		if err != nil {
			return nil, err
		}

		if !rbcid.Equals(k) {
			return nil, ErrHashMismatch
		}

		return blocks.NewBlockWithCid(bdata, rbcid)
	} else {
		return blocks.NewBlockWithCid(bdata, k)
	}
}

func (bs *blockstore) Put(block blocks.Block) error {
	k := dshelp.CidToDsKey(block.Cid())

	// Has is cheaper than Put, so see if we already have it
	exists, err := bs.datastore.Has(k)
	if err == nil && exists {
		return nil // already stored.
	}
	return bs.datastore.Put(k, block.RawData())
}

func (bs *blockstore) PutMany(blocks []blocks.Block) error {
	t, err := bs.datastore.Batch()
	if err != nil {
		return err
	}
	for _, b := range blocks {
		k := dshelp.CidToDsKey(b.Cid())
		exists, err := bs.datastore.Has(k)
		if err == nil && exists {
			continue
		}

		err = t.Put(k, b.RawData())
		if err != nil {
			return err
		}
	}
	return t.Commit()
}

func (bs *blockstore) Has(k *cid.Cid) (bool, error) {
	return bs.datastore.Has(dshelp.CidToDsKey(k))
}

func (s *blockstore) DeleteBlock(k *cid.Cid) error {
	return s.datastore.Delete(dshelp.CidToDsKey(k))
}

// AllKeysChan runs a query for keys from the blockstore.
// this is very simplistic, in the future, take dsq.Query as a param?
//
// AllKeysChan respects context
func (bs *blockstore) AllKeysChan(ctx context.Context) (<-chan *cid.Cid, error) {

	// KeysOnly, because that would be _a lot_ of data.
	q := dsq.Query{KeysOnly: true}
	// datastore/namespace does *NOT* fix up Query.Prefix
	q.Prefix = BlockPrefix.String()
	res, err := bs.datastore.Query(q)
	if err != nil {
		return nil, err
	}

	output := make(chan *cid.Cid, dsq.KeysOnlyBufSize)
	go func() {
		defer func() {
			res.Close() // ensure exit (signals early exit, too)
			close(output)
		}()

		for {
			e, ok := res.NextSync()
			if !ok {
				return
			}
			if e.Error != nil {
				log.Debug("blockstore.AllKeysChan got err:", e.Error)
				continue
			}

			// need to convert to key.Key using key.KeyFromDsKey.
			k, err := dshelp.DsKeyToCid(ds.RawKey(e.Key))
			if err != nil {
				log.Warningf("error parsing key from DsKey: ", err)
				continue
			}

			select {
			case <-ctx.Done():
				return
			case output <- k:
			}
		}
	}()

	return output, nil
}

func NewGCLocker() *gclocker {
	return &gclocker{}
}

type gclocker struct {
	lk      sync.RWMutex
	gcreq   int32
	gcreqlk sync.Mutex
}

type Unlocker interface {
	Unlock()
}

type unlocker struct {
	unlock func()
}

func (u *unlocker) Unlock() {
	u.unlock()
	u.unlock = nil // ensure its not called twice
}

func (bs *gclocker) GCLock() Unlocker {
	atomic.AddInt32(&bs.gcreq, 1)
	bs.lk.Lock()
	atomic.AddInt32(&bs.gcreq, -1)
	return &unlocker{bs.lk.Unlock}
}

func (bs *gclocker) PinLock() Unlocker {
	bs.lk.RLock()
	return &unlocker{bs.lk.RUnlock}
}

func (bs *gclocker) GCRequested() bool {
	return atomic.LoadInt32(&bs.gcreq) > 0
}
<commit_msg>blockstore.AllKeyChan: fix/cleanup error handling<commit_after>// package blockstore implements a thin wrapper over a datastore, giving a
// clean interface for Getting and Putting block objects.
package blockstore

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"

	blocks "github.com/ipfs/go-ipfs/blocks"
	dshelp "github.com/ipfs/go-ipfs/thirdparty/ds-help"

	ds "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore"
	dsns "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore/namespace"
	dsq "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore/query"
	logging "gx/ipfs/QmSpJByNKFX1sCsHBEp3R73FL4NF6FnQTEGyNAXHm2GS52/go-log"
	cid "gx/ipfs/QmcTcsTvfaeEBRFo1TkFgT8sRmgi1n1LTZpecfVP8fzpGD/go-cid"
)

var log = logging.Logger("blockstore")

// BlockPrefix namespaces blockstore datastores
var BlockPrefix = ds.NewKey("blocks")

var ValueTypeMismatch = errors.New("the retrieved value is not a Block")
var ErrHashMismatch = errors.New("block in storage has different hash than requested")

var ErrNotFound = errors.New("blockstore: block not found")

// Blockstore wraps a Datastore
type Blockstore interface {
	DeleteBlock(*cid.Cid) error
	Has(*cid.Cid) (bool, error)
	Get(*cid.Cid) (blocks.Block, error)
	Put(blocks.Block) error
	PutMany([]blocks.Block) error

	AllKeysChan(ctx context.Context) (<-chan *cid.Cid, error)
}

type GCLocker interface {
	// GCLock locks the blockstore for garbage collection. No operations
	// that expect to finish with a pin should ocurr simultaneously.
	// Reading during GC is safe, and requires no lock.
	GCLock() Unlocker

	// PinLock locks the blockstore for sequences of puts expected to finish
	// with a pin (before GC). Multiple put->pin sequences can write through
	// at the same time, but no GC should not happen simulatenously.
	// Reading during Pinning is safe, and requires no lock.
	PinLock() Unlocker

	// GcRequested returns true if GCLock has been called and is waiting to
	// take the lock
	GCRequested() bool
}

type GCBlockstore interface {
	Blockstore
	GCLocker
}

func NewGCBlockstore(bs Blockstore, gcl GCLocker) GCBlockstore {
	return gcBlockstore{bs, gcl}
}

type gcBlockstore struct {
	Blockstore
	GCLocker
}

func NewBlockstore(d ds.Batching) *blockstore {
	var dsb ds.Batching
	dd := dsns.Wrap(d, BlockPrefix)
	dsb = dd
	return &blockstore{
		datastore: dsb,
	}
}

type blockstore struct {
	datastore ds.Batching

	lk      sync.RWMutex
	gcreq   int32
	gcreqlk sync.Mutex

	rehash bool
}

func (bs *blockstore) HashOnRead(enabled bool) {
	bs.rehash = enabled
}

func (bs *blockstore) Get(k *cid.Cid) (blocks.Block, error) {
	if k == nil {
		log.Error("nil cid in blockstore")
		return nil, ErrNotFound
	}

	maybeData, err := bs.datastore.Get(dshelp.CidToDsKey(k))
	if err == ds.ErrNotFound {
		return nil, ErrNotFound
	}
	if err != nil {
		return nil, err
	}
	bdata, ok := maybeData.([]byte)
	if !ok {
		return nil, ValueTypeMismatch
	}

	if bs.rehash {
		rbcid, err := k.Prefix().Sum(bdata)
		if err != nil {
			return nil, err
		}

		if !rbcid.Equals(k) {
			return nil, ErrHashMismatch
		}

		return blocks.NewBlockWithCid(bdata, rbcid)
	} else {
		return blocks.NewBlockWithCid(bdata, k)
	}
}

func (bs *blockstore) Put(block blocks.Block) error {
	k := dshelp.CidToDsKey(block.Cid())

	// Has is cheaper than Put, so see if we already have it
	exists, err := bs.datastore.Has(k)
	if err == nil && exists {
		return nil // already stored.
	}
	return bs.datastore.Put(k, block.RawData())
}

func (bs *blockstore) PutMany(blocks []blocks.Block) error {
	t, err := bs.datastore.Batch()
	if err != nil {
		return err
	}
	for _, b := range blocks {
		k := dshelp.CidToDsKey(b.Cid())
		exists, err := bs.datastore.Has(k)
		if err == nil && exists {
			continue
		}

		err = t.Put(k, b.RawData())
		if err != nil {
			return err
		}
	}
	return t.Commit()
}

func (bs *blockstore) Has(k *cid.Cid) (bool, error) {
	return bs.datastore.Has(dshelp.CidToDsKey(k))
}

func (s *blockstore) DeleteBlock(k *cid.Cid) error {
	return s.datastore.Delete(dshelp.CidToDsKey(k))
}

// AllKeysChan runs a query for keys from the blockstore.
// this is very simplistic, in the future, take dsq.Query as a param?
//
// AllKeysChan respects context
func (bs *blockstore) AllKeysChan(ctx context.Context) (<-chan *cid.Cid, error) {

	// KeysOnly, because that would be _a lot_ of data.
	q := dsq.Query{KeysOnly: true}
	// datastore/namespace does *NOT* fix up Query.Prefix
	q.Prefix = BlockPrefix.String()
	res, err := bs.datastore.Query(q)
	if err != nil {
		return nil, err
	}

	output := make(chan *cid.Cid, dsq.KeysOnlyBufSize)
	go func() {
		defer func() {
			res.Close() // ensure exit (signals early exit, too)
			close(output)
		}()

		for {
			e, ok := res.NextSync()
			if !ok {
				return
			}
			if e.Error != nil {
				log.Errorf("blockstore.AllKeysChan got err:", e.Error)
				return
			}

			// need to convert to key.Key using key.KeyFromDsKey.
			k, err := dshelp.DsKeyToCid(ds.RawKey(e.Key))
			if err != nil {
				log.Warningf("error parsing key from DsKey: ", err)
				continue
			}

			select {
			case <-ctx.Done():
				return
			case output <- k:
			}
		}
	}()

	return output, nil
}

func NewGCLocker() *gclocker {
	return &gclocker{}
}

type gclocker struct {
	lk      sync.RWMutex
	gcreq   int32
	gcreqlk sync.Mutex
}

type Unlocker interface {
	Unlock()
}

type unlocker struct {
	unlock func()
}

func (u *unlocker) Unlock() {
	u.unlock()
	u.unlock = nil // ensure its not called twice
}

func (bs *gclocker) GCLock() Unlocker {
	atomic.AddInt32(&bs.gcreq, 1)
	bs.lk.Lock()
	atomic.AddInt32(&bs.gcreq, -1)
	return &unlocker{bs.lk.Unlock}
}

func (bs *gclocker) PinLock() Unlocker {
	bs.lk.RLock()
	return &unlocker{bs.lk.RUnlock}
}

func (bs *gclocker) GCRequested() bool {
	return atomic.LoadInt32(&bs.gcreq) > 0
}
<|endoftext|>


<commit_before>/**
 * MediaType
 *
 * Copyright © 2014 Trevor N. Suarez (Rican7)
 */

package mediatype

import (
	"testing"
)

/**
 * Variables
 */

var validComplexMediaType = "application/vnd.google-earth.kml+xml; charset=utf-8"

var validMediaTypeStrings = []string{
	"example",
	"application/json",
	"application/xhtml+xml",
	"audio/vnd.rn-realaudio",
	"image/vnd.djvu",
}

var invalidMediaTypeStrings = []string{
	"application/vnd/json",
	"invalid/+json.really/sick",
}

/**
 * Helper functions
 */

// Parse our valid media types
func parseValidMediaTypes() (map[string]MediaType, map[string]error) {
	mediaTypes := make(map[string]MediaType)
	errors := make(map[string]error)

	for _, val := range validMediaTypeStrings {
		mt, err := Parse(val)

		if nil != mt {
			mediaTypes[val] = mt
		}

		if nil != err {
			errors[val] = err
		}
	}

	return mediaTypes, errors
}

// Parse our invalid media types
func parseInvalidMediaTypes() (map[string]MediaType, map[string]error) {
	mediaTypes := make(map[string]MediaType)
	errors := make(map[string]error)

	for _, val := range invalidMediaTypeStrings {
		mt, err := Parse(val)

		if nil != mt {
			mediaTypes[val] = mt
		}

		if nil != err {
			errors[val] = err
		}
	}

	return mediaTypes, errors
}

/**
 * Tests functions
 */

func TestParse(t *testing.T) {
	_, validErrs := parseValidMediaTypes()
	invalid, _ := parseInvalidMediaTypes()

	for key, err := range validErrs {
		t.Errorf("Parsing failed for valid '%s' with error '%s'", key, err)
	}

	for key, _ := range invalid {
		t.Errorf("Parsing succeeded for invalid '%s'", key)
	}
}

func TestFullType(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.FullType() != "application/vnd.google-earth.kml+xml" {
			t.Errorf("Incorrect full type for %+v", mt)
		}
	}
}

func TestParameters(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {
		correctParameters := map[string]string{"charset": "utf-8"}

		for i, tree := range mt.Parameters() {
			if tree != correctParameters[i] {
				t.Errorf("Incorrect parameters for %+v", mt)
			}
		}
	}
}

func TestMainType(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.MainType() != "application" {
			t.Errorf("Incorrect main type for %+v", mt)
		}
	}
}

func TestSubType(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.SubType() != "kml" {
			t.Errorf("Incorrect sub type for %+v", mt)
		}
	}
}

func TestTrees(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {
		correctTrees := []string{"vnd", "google-earth"}

		for i, tree := range mt.Trees() {
			if tree != correctTrees[i] {
				t.Errorf("Incorrect trees for %+v", mt)
			}
		}
	}
}

func TestPrefix(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.Prefix() != "vnd" {
			t.Errorf("Incorrect prefix for %+v", mt)
		}
	}
}

func TestSuffix(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.Suffix() != "xml" {
			t.Errorf("Incorrect suffix for %+v", mt)
		}
	}
}

func TestString(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.String() != validComplexMediaType {
			t.Errorf("Incorrect string for %+v", mt)
		}
	}
}
<commit_msg>Oh... uhhhh, thanks golint!<commit_after>/**
 * MediaType
 *
 * Copyright © 2014 Trevor N. Suarez (Rican7)
 */

package mediatype

import (
	"testing"
)

/**
 * Variables
 */

var validComplexMediaType = "application/vnd.google-earth.kml+xml; charset=utf-8"

var validMediaTypeStrings = []string{
	"example",
	"application/json",
	"application/xhtml+xml",
	"audio/vnd.rn-realaudio",
	"image/vnd.djvu",
}

var invalidMediaTypeStrings = []string{
	"application/vnd/json",
	"invalid/+json.really/sick",
}

/**
 * Helper functions
 */

// Parse our valid media types
func parseValidMediaTypes() (map[string]MediaType, map[string]error) {
	mediaTypes := make(map[string]MediaType)
	errors := make(map[string]error)

	for _, val := range validMediaTypeStrings {
		mt, err := Parse(val)

		if nil != mt {
			mediaTypes[val] = mt
		}

		if nil != err {
			errors[val] = err
		}
	}

	return mediaTypes, errors
}

// Parse our invalid media types
func parseInvalidMediaTypes() (map[string]MediaType, map[string]error) {
	mediaTypes := make(map[string]MediaType)
	errors := make(map[string]error)

	for _, val := range invalidMediaTypeStrings {
		mt, err := Parse(val)

		if nil != mt {
			mediaTypes[val] = mt
		}

		if nil != err {
			errors[val] = err
		}
	}

	return mediaTypes, errors
}

/**
 * Tests functions
 */

func TestParse(t *testing.T) {
	_, validErrs := parseValidMediaTypes()
	invalid, _ := parseInvalidMediaTypes()

	for key, err := range validErrs {
		t.Errorf("Parsing failed for valid '%s' with error '%s'", key, err)
	}

	for key := range invalid {
		t.Errorf("Parsing succeeded for invalid '%s'", key)
	}
}

func TestFullType(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.FullType() != "application/vnd.google-earth.kml+xml" {
			t.Errorf("Incorrect full type for %+v", mt)
		}
	}
}

func TestParameters(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {
		correctParameters := map[string]string{"charset": "utf-8"}

		for i, tree := range mt.Parameters() {
			if tree != correctParameters[i] {
				t.Errorf("Incorrect parameters for %+v", mt)
			}
		}
	}
}

func TestMainType(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.MainType() != "application" {
			t.Errorf("Incorrect main type for %+v", mt)
		}
	}
}

func TestSubType(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.SubType() != "kml" {
			t.Errorf("Incorrect sub type for %+v", mt)
		}
	}
}

func TestTrees(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {
		correctTrees := []string{"vnd", "google-earth"}

		for i, tree := range mt.Trees() {
			if tree != correctTrees[i] {
				t.Errorf("Incorrect trees for %+v", mt)
			}
		}
	}
}

func TestPrefix(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.Prefix() != "vnd" {
			t.Errorf("Incorrect prefix for %+v", mt)
		}
	}
}

func TestSuffix(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.Suffix() != "xml" {
			t.Errorf("Incorrect suffix for %+v", mt)
		}
	}
}

func TestString(t *testing.T) {
	mt, err := Parse(validComplexMediaType)

	if nil != err {
		t.Errorf("Parsing failed for valid '%s'", validComplexMediaType)
	} else {

		if mt.String() != validComplexMediaType {
			t.Errorf("Incorrect string for %+v", mt)
		}
	}
}
<|endoftext|>


<commit_before>// Based on ssh/terminal:
// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build appengine

package logrus

func initTerminal(w io.Writer) {
}
<commit_msg>Fix copypasta<commit_after>// Based on ssh/terminal:
// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build appengine

package logrus

import "io"

func initTerminal(w io.Writer) {
}
<|endoftext|>


<commit_before>package cache

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"reflect"
	"sync/atomic"
	"time"

	"github.com/VictoriaMetrics/fastcache"
	"github.com/go-redis/redis/v8"
	"github.com/klauspost/compress/s2"
	"github.com/vmihailenco/bufpool"
	"github.com/vmihailenco/msgpack/v5"
	"go4.org/syncutil/singleflight"
)

const compressionThreshold = 64

const (
	noCompression = 0x0
	s2Compression = 0x1
)

var ErrCacheMiss = errors.New("cache: key is missing")
var errRedisLocalCacheNil = errors.New("cache: both Redis and LocalCache are nil")

type rediser interface {
	Set(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.StatusCmd
	SetXX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd
	SetNX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd

	Get(ctx context.Context, key string) *redis.StringCmd
	Del(ctx context.Context, keys ...string) *redis.IntCmd
}

type Item struct {
	Ctx context.Context

	Key   string
	Value interface{}

	// TTL is the cache expiration time.
	// Default TTL is 1 hour.
	TTL time.Duration

	// Do returns value to be cached.
	Do func(*Item) (interface{}, error)

	// IfExists only sets the key if it already exist.
	IfExists bool

	// IfNotExists only sets the key if it does not already exist.
	IfNotExists bool

	// SkipLocalCache skips local cache as if it is not set.
	SkipLocalCache bool
}

func (item *Item) Context() context.Context {
	if item.Ctx == nil {
		return context.Background()
	}
	return item.Ctx
}

func (item *Item) value() (interface{}, error) {
	if item.Do != nil {
		return item.Do(item)
	}
	if item.Value != nil {
		return item.Value, nil
	}
	return nil, nil
}

func (item *Item) ttl() time.Duration {
	if item.TTL < 0 {
		return 0
	}
	if item.TTL < time.Second {
		return time.Hour
	}
	return item.TTL
}

//------------------------------------------------------------------------------

type Options struct {
	Redis rediser

	LocalCache    *fastcache.Cache
	LocalCacheTTL time.Duration

	StatsEnabled bool
}

func (opt *Options) init() {
	switch opt.LocalCacheTTL {
	case -1:
		opt.LocalCacheTTL = 0
	case 0:
		opt.LocalCacheTTL = time.Minute
	}
}

type Cache struct {
	opt *Options

	group singleflight.Group

	hits   uint64
	misses uint64
}

func New(opt *Options) *Cache {
	opt.init()
	return &Cache{
		opt: opt,
	}
}

// Set caches the item.
func (cd *Cache) Set(item *Item) error {
	_, _, err := cd.set(item)
	return err
}

func (cd *Cache) set(item *Item) ([]byte, bool, error) {
	value, err := item.value()
	if err != nil {
		return nil, false, err
	}

	b, err := cd.Marshal(value)
	if err != nil {
		return nil, false, err
	}

	if cd.opt.LocalCache != nil {
		cd.localSet(item.Key, b)
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return b, true, errRedisLocalCacheNil
		}
		return b, true, nil
	}

	if item.IfExists {
		return b, true, cd.opt.Redis.SetXX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	if item.IfNotExists {
		return b, true, cd.opt.Redis.SetNX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	return b, true, cd.opt.Redis.Set(item.Context(), item.Key, b, item.ttl()).Err()
}

// Exists reports whether value for the given key exists.
func (cd *Cache) Exists(ctx context.Context, key string) bool {
	return cd.Get(ctx, key, nil) == nil
}

// Get gets the value for the given key.
func (cd *Cache) Get(ctx context.Context, key string, value interface{}) error {
	return cd.get(ctx, key, value, false)
}

// Get gets the value for the given key skipping local cache.
func (cd *Cache) GetSkippingLocalCache(
	ctx context.Context, key string, value interface{},
) error {
	return cd.get(ctx, key, value, true)
}

func (cd *Cache) get(
	ctx context.Context,
	key string,
	value interface{},
	skipLocalCache bool,
) error {
	b, err := cd.getBytes(ctx, key, skipLocalCache)
	if err != nil {
		return err
	}
	return cd.Unmarshal(b, value)
}

func (cd *Cache) getBytes(ctx context.Context, key string, skipLocalCache bool) ([]byte, error) {
	if !skipLocalCache && cd.opt.LocalCache != nil {
		b, ok := cd.localGet(key)
		if ok {
			return b, nil
		}
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return nil, errRedisLocalCacheNil
		}
		return nil, ErrCacheMiss
	}

	b, err := cd.opt.Redis.Get(ctx, key).Bytes()
	if err != nil {
		if cd.opt.StatsEnabled {
			atomic.AddUint64(&cd.misses, 1)
		}
		if err == redis.Nil {
			return nil, ErrCacheMiss
		}
		return nil, err
	}

	if cd.opt.StatsEnabled {
		atomic.AddUint64(&cd.hits, 1)
	}

	if !skipLocalCache && cd.opt.LocalCache != nil {
		cd.localSet(key, b)
	}
	return b, nil
}

// Once gets the item.Value for the given item.Key from the cache or
// executes, caches, and returns the results of the given item.Func,
// making sure that only one execution is in-flight for a given item.Key
// at a time. If a duplicate comes in, the duplicate caller waits for the
// original to complete and receives the same results.
func (cd *Cache) Once(item *Item) error {
	b, cached, err := cd.getSetItemBytesOnce(item)
	if err != nil {
		return err
	}

	if item.Value == nil || len(b) == 0 {
		return nil
	}

	if err := cd.Unmarshal(b, item.Value); err != nil {
		if cached {
			_ = cd.Delete(item.Context(), item.Key)
			return cd.Once(item)
		}
		return err
	}

	return nil
}

func (cd *Cache) getSetItemBytesOnce(item *Item) (b []byte, cached bool, err error) {
	if cd.opt.LocalCache != nil {
		b, ok := cd.localGet(item.Key)
		if ok {
			return b, true, nil
		}
	}

	v, err := cd.group.Do(item.Key, func() (interface{}, error) {
		b, err := cd.getBytes(item.Context(), item.Key, item.SkipLocalCache)
		if err == nil {
			cached = true
			return b, nil
		}

		b, ok, err := cd.set(item)
		if ok {
			return b, nil
		}
		return nil, err
	})
	if err != nil {
		return nil, false, err
	}
	return v.([]byte), cached, nil
}

func (cd *Cache) Delete(ctx context.Context, key string) error {
	if cd.opt.LocalCache != nil {
		cd.opt.LocalCache.Del([]byte(key))
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return errRedisLocalCacheNil
		}
		return nil
	}

	deleted, err := cd.opt.Redis.Del(ctx, key).Result()
	if err != nil {
		return err
	}
	if deleted == 0 {
		return ErrCacheMiss
	}
	return nil
}

func (cd *Cache) localSet(key string, b []byte) {
	if cd.opt.LocalCacheTTL > 0 {
		pos := len(b)
		b = append(b, make([]byte, 4)...)
		encodeTime(b[pos:], time.Now())
	}

	cd.opt.LocalCache.Set([]byte(key), b)
}

func (cd *Cache) localGet(key string) ([]byte, bool) {
	b, ok := cd.opt.LocalCache.HasGet(nil, []byte(key))
	if !ok {
		return b, false
	}

	if len(b) == 0 || cd.opt.LocalCacheTTL == 0 {
		return b, true
	}
	if len(b) < 4 {
		panic("not reached")
	}

	tm := decodeTime(b[len(b)-4:])
	if time.Since(tm) > cd.opt.LocalCacheTTL {
		cd.opt.LocalCache.Del([]byte(key))
		return nil, false
	}

	return b[:len(b)-4], true
}

func (cd *Cache) Marshal(value interface{}) ([]byte, error) {
	switch value := value.(type) {
	case nil:
		return nil, nil
	case []byte:
		return value, nil
	case string:
		return []byte(value), nil
	}

	enc := msgpack.GetEncoder()

	var buf bytes.Buffer
	enc.Reset(&buf)
	enc.UseCompactInts(true)

	err := enc.Encode(value)

	msgpack.PutEncoder(enc)

	if err != nil {
		return nil, err
	}

	b := buf.Bytes()

	if len(b) < compressionThreshold {
		b = append(b, noCompression)
		return b, nil
	}

	b = s2.Encode(nil, b)
	b = append(b, s2Compression)

	return b, nil
}

func (cd *Cache) Unmarshal(b []byte, value interface{}) error {
	if len(b) == 0 {
		return nil
	}

	switch value := value.(type) {
	case nil:
		return nil
	case *[]byte:
		reflect.ValueOf(value).Elem().SetBytes(b)
		return nil
	case *string:
		reflect.ValueOf(value).Elem().SetString(string(b))
		return nil
	}

	if len(b) == 0 {
		return nil
	}

	switch c := b[len(b)-1]; c {
	case noCompression:
		b = b[:len(b)-1]
	case s2Compression:
		b = b[:len(b)-1]

		n, err := s2.DecodedLen(b)
		if err != nil {
			return err
		}

		buf := bufpool.Get(n)
		defer bufpool.Put(buf)

		b, err = s2.Decode(buf.Bytes(), b)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("uknownn compression method: %x", c)
	}

	return msgpack.Unmarshal(b, value)
}

//------------------------------------------------------------------------------

type Stats struct {
	Hits   uint64
	Misses uint64
}

// Stats returns cache statistics.
func (cd *Cache) Stats() *Stats {
	if !cd.opt.StatsEnabled {
		return nil
	}
	return &Stats{
		Hits:   atomic.LoadUint64(&cd.hits),
		Misses: atomic.LoadUint64(&cd.misses),
	}
}

//------------------------------------------------------------------------------

var epoch = time.Date(2020, time.January, 01, 00, 0, 0, 0, time.UTC).Unix()

func encodeTime(b []byte, tm time.Time) {
	secs := tm.Unix() - epoch
	binary.LittleEndian.PutUint32(b, uint32(secs))
}

func decodeTime(b []byte) time.Time {
	secs := binary.LittleEndian.Uint32(b)
	return time.Unix(int64(secs)+epoch, 0)
}
<commit_msg>Reduce number of allocs<commit_after>package cache

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"sync/atomic"
	"time"

	"github.com/VictoriaMetrics/fastcache"
	"github.com/go-redis/redis/v8"
	"github.com/klauspost/compress/s2"
	"github.com/vmihailenco/bufpool"
	"github.com/vmihailenco/msgpack/v5"
	"go4.org/syncutil/singleflight"
)

const compressionThreshold = 64

const (
	noCompression = 0x0
	s2Compression = 0x1
)

var ErrCacheMiss = errors.New("cache: key is missing")
var errRedisLocalCacheNil = errors.New("cache: both Redis and LocalCache are nil")

type rediser interface {
	Set(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.StatusCmd
	SetXX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd
	SetNX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd

	Get(ctx context.Context, key string) *redis.StringCmd
	Del(ctx context.Context, keys ...string) *redis.IntCmd
}

type Item struct {
	Ctx context.Context

	Key   string
	Value interface{}

	// TTL is the cache expiration time.
	// Default TTL is 1 hour.
	TTL time.Duration

	// Do returns value to be cached.
	Do func(*Item) (interface{}, error)

	// IfExists only sets the key if it already exist.
	IfExists bool

	// IfNotExists only sets the key if it does not already exist.
	IfNotExists bool

	// SkipLocalCache skips local cache as if it is not set.
	SkipLocalCache bool
}

func (item *Item) Context() context.Context {
	if item.Ctx == nil {
		return context.Background()
	}
	return item.Ctx
}

func (item *Item) value() (interface{}, error) {
	if item.Do != nil {
		return item.Do(item)
	}
	if item.Value != nil {
		return item.Value, nil
	}
	return nil, nil
}

func (item *Item) ttl() time.Duration {
	if item.TTL < 0 {
		return 0
	}
	if item.TTL < time.Second {
		return time.Hour
	}
	return item.TTL
}

//------------------------------------------------------------------------------

type Options struct {
	Redis rediser

	LocalCache    *fastcache.Cache
	LocalCacheTTL time.Duration

	StatsEnabled bool
}

func (opt *Options) init() {
	switch opt.LocalCacheTTL {
	case -1:
		opt.LocalCacheTTL = 0
	case 0:
		opt.LocalCacheTTL = time.Minute
	}
}

type Cache struct {
	opt *Options

	group   singleflight.Group
	bufpool bufpool.Pool

	hits   uint64
	misses uint64
}

func New(opt *Options) *Cache {
	opt.init()
	return &Cache{
		opt: opt,
	}
}

// Set caches the item.
func (cd *Cache) Set(item *Item) error {
	_, _, err := cd.set(item)
	return err
}

func (cd *Cache) set(item *Item) ([]byte, bool, error) {
	value, err := item.value()
	if err != nil {
		return nil, false, err
	}

	b, err := cd.Marshal(value)
	if err != nil {
		return nil, false, err
	}

	if cd.opt.LocalCache != nil {
		cd.localSet(item.Key, b)
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return b, true, errRedisLocalCacheNil
		}
		return b, true, nil
	}

	if item.IfExists {
		return b, true, cd.opt.Redis.SetXX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	if item.IfNotExists {
		return b, true, cd.opt.Redis.SetNX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	return b, true, cd.opt.Redis.Set(item.Context(), item.Key, b, item.ttl()).Err()
}

// Exists reports whether value for the given key exists.
func (cd *Cache) Exists(ctx context.Context, key string) bool {
	return cd.Get(ctx, key, nil) == nil
}

// Get gets the value for the given key.
func (cd *Cache) Get(ctx context.Context, key string, value interface{}) error {
	return cd.get(ctx, key, value, false)
}

// Get gets the value for the given key skipping local cache.
func (cd *Cache) GetSkippingLocalCache(
	ctx context.Context, key string, value interface{},
) error {
	return cd.get(ctx, key, value, true)
}

func (cd *Cache) get(
	ctx context.Context,
	key string,
	value interface{},
	skipLocalCache bool,
) error {
	b, err := cd.getBytes(ctx, key, skipLocalCache)
	if err != nil {
		return err
	}
	return cd.Unmarshal(b, value)
}

func (cd *Cache) getBytes(ctx context.Context, key string, skipLocalCache bool) ([]byte, error) {
	if !skipLocalCache && cd.opt.LocalCache != nil {
		b, ok := cd.localGet(key)
		if ok {
			return b, nil
		}
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return nil, errRedisLocalCacheNil
		}
		return nil, ErrCacheMiss
	}

	b, err := cd.opt.Redis.Get(ctx, key).Bytes()
	if err != nil {
		if cd.opt.StatsEnabled {
			atomic.AddUint64(&cd.misses, 1)
		}
		if err == redis.Nil {
			return nil, ErrCacheMiss
		}
		return nil, err
	}

	if cd.opt.StatsEnabled {
		atomic.AddUint64(&cd.hits, 1)
	}

	if !skipLocalCache && cd.opt.LocalCache != nil {
		cd.localSet(key, b)
	}
	return b, nil
}

// Once gets the item.Value for the given item.Key from the cache or
// executes, caches, and returns the results of the given item.Func,
// making sure that only one execution is in-flight for a given item.Key
// at a time. If a duplicate comes in, the duplicate caller waits for the
// original to complete and receives the same results.
func (cd *Cache) Once(item *Item) error {
	b, cached, err := cd.getSetItemBytesOnce(item)
	if err != nil {
		return err
	}

	if item.Value == nil || len(b) == 0 {
		return nil
	}

	if err := cd.Unmarshal(b, item.Value); err != nil {
		if cached {
			_ = cd.Delete(item.Context(), item.Key)
			return cd.Once(item)
		}
		return err
	}

	return nil
}

func (cd *Cache) getSetItemBytesOnce(item *Item) (b []byte, cached bool, err error) {
	if cd.opt.LocalCache != nil {
		b, ok := cd.localGet(item.Key)
		if ok {
			return b, true, nil
		}
	}

	v, err := cd.group.Do(item.Key, func() (interface{}, error) {
		b, err := cd.getBytes(item.Context(), item.Key, item.SkipLocalCache)
		if err == nil {
			cached = true
			return b, nil
		}

		b, ok, err := cd.set(item)
		if ok {
			return b, nil
		}
		return nil, err
	})
	if err != nil {
		return nil, false, err
	}
	return v.([]byte), cached, nil
}

func (cd *Cache) Delete(ctx context.Context, key string) error {
	if cd.opt.LocalCache != nil {
		cd.opt.LocalCache.Del([]byte(key))
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return errRedisLocalCacheNil
		}
		return nil
	}

	deleted, err := cd.opt.Redis.Del(ctx, key).Result()
	if err != nil {
		return err
	}
	if deleted == 0 {
		return ErrCacheMiss
	}
	return nil
}

func (cd *Cache) localSet(key string, b []byte) {
	if cd.opt.LocalCacheTTL > 0 {
		pos := len(b)
		b = append(b, make([]byte, 4)...)
		encodeTime(b[pos:], time.Now())
	}

	cd.opt.LocalCache.Set([]byte(key), b)
}

func (cd *Cache) localGet(key string) ([]byte, bool) {
	b, ok := cd.opt.LocalCache.HasGet(nil, []byte(key))
	if !ok {
		return b, false
	}

	if len(b) == 0 || cd.opt.LocalCacheTTL == 0 {
		return b, true
	}
	if len(b) < 4 {
		panic("not reached")
	}

	tm := decodeTime(b[len(b)-4:])
	if time.Since(tm) > cd.opt.LocalCacheTTL {
		cd.opt.LocalCache.Del([]byte(key))
		return nil, false
	}

	return b[:len(b)-4], true
}

func (cd *Cache) Marshal(value interface{}) ([]byte, error) {
	switch value := value.(type) {
	case nil:
		return nil, nil
	case []byte:
		return value, nil
	case string:
		return []byte(value), nil
	}

	buf := cd.bufpool.Get()
	defer cd.bufpool.Put(buf)

	enc := msgpack.GetEncoder()
	enc.Reset(buf)
	enc.UseCompactInts(true)

	err := enc.Encode(value)

	msgpack.PutEncoder(enc)

	if err != nil {
		return nil, err
	}

	if buf.Len() < compressionThreshold {
		b := make([]byte, buf.Len()+1)
		copy(b, buf.Bytes())
		b[len(b)-1] = noCompression
		return b, nil
	}

	b := make([]byte, s2.MaxEncodedLen(buf.Len())+1)
	b = s2.Encode(b, buf.Bytes())
	b = append(b, s2Compression)

	return b, nil
}

func (cd *Cache) Unmarshal(b []byte, value interface{}) error {
	if len(b) == 0 {
		return nil
	}

	switch value := value.(type) {
	case nil:
		return nil
	case *[]byte:
		*value = b
		return nil
	case *string:
		*value = string(b)
		return nil
	}

	if len(b) == 0 {
		return nil
	}

	switch c := b[len(b)-1]; c {
	case noCompression:
		b = b[:len(b)-1]
	case s2Compression:
		b = b[:len(b)-1]

		n, err := s2.DecodedLen(b)
		if err != nil {
			return err
		}

		buf := bufpool.Get(n)
		defer bufpool.Put(buf)

		b, err = s2.Decode(buf.Bytes(), b)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("uknownn compression method: %x", c)
	}

	return msgpack.Unmarshal(b, value)
}

//------------------------------------------------------------------------------

type Stats struct {
	Hits   uint64
	Misses uint64
}

// Stats returns cache statistics.
func (cd *Cache) Stats() *Stats {
	if !cd.opt.StatsEnabled {
		return nil
	}
	return &Stats{
		Hits:   atomic.LoadUint64(&cd.hits),
		Misses: atomic.LoadUint64(&cd.misses),
	}
}

//------------------------------------------------------------------------------

var epoch = time.Date(2020, time.January, 01, 00, 0, 0, 0, time.UTC).Unix()

func encodeTime(b []byte, tm time.Time) {
	secs := tm.Unix() - epoch
	binary.LittleEndian.PutUint32(b, uint32(secs))
}

func decodeTime(b []byte) time.Time {
	secs := binary.LittleEndian.Uint32(b)
	return time.Unix(int64(secs)+epoch, 0)
}
<|endoftext|>