<commit_before>package isumm

import (
	"fmt"
	"reflect"
	"sort"
	"testing"
	"time"
)

var (
	now       = time.Now()
	nextMonth = now.Add(32 * 24 * time.Hour)
	incMin    = now.Add(10 * time.Minute)
	opType    = "1"
	value     = "12"
	date      = "2006-01-02"
)

func TestSort(t *testing.T) {
	o := Operations{NewOperation(Balance, 1, now), NewOperation(Balance, 1, incMin)}
	sort.Sort(o)
	if o[0].Date().Before(o[1].Date()) {
		t.Errorf("want %s before than %s", o[0], o[1])
	}
}

func TestNewOperationFromString(t *testing.T) {
	validOp, _ := NewOperationFromString(opType, value, date)
	testCases := []struct {
		desc, t, v, d string
		want          *Operation
	}{
		// Invalid cases.
		{desc: "invalid value - empty", t: opType, v: "", d: date},
		{desc: "invalid value - chars", t: opType, v: "acb", d: date},
		{desc: "invalid op - empty string", t: "", v: value, d: date},
		{desc: "invalid op - type does not exist", t: "94879138ddffg", v: value, d: date},
		{desc: "invalid date", t: opType, v: value, d: "31/31/31"},
		// Valid cases.
		{desc: "valid - contain spaces", t: opType, v: fmt.Sprintf(" %s ", value), d: date, want: &validOp},
		{desc: "valid - perfect", t: opType, v: value, d: date, want: &validOp},
	}
	for _, test := range testCases {
		got, err := NewOperationFromString(test.t, test.v, test.d)
		switch {
		case test.want == nil:
			if err == nil {
				t.Errorf("got:nil expected:err. Test:(%+v)", test)
			}
		default:
			if !reflect.DeepEqual(test.want, &got) {
				t.Errorf("got:(%+v) want:(%+v). Test:(%+v)", got, test.want, test)
			}
		}
	}
}

func TestSummarize(t *testing.T) {
	data := []struct {
		ops  Operations
		summ Summaries
	}{
		{ // Simple case: One entry containing only balance.
			ops:  Operations{NewOperation(Balance, 1.2, now)},
			summ: Summaries{{Date: monthYear(now), Balance: 1.2, Change: 0}},
		},
		{ // Empty case.
			ops:  Operations{},
			summ: Summaries{},
		},
		{
			ops: Operations{
				NewOperation(Balance, 1.2, incMin),
				NewOperation(Deposit, 1.0, now),
				NewOperation(Balance, 2.2, nextMonth),
			},
			summ: Summaries{
				{Date: monthYear(nextMonth), Balance: 2.2, Change: 0},
				{Date: monthYear(now), Balance: 1.2, Change: 1.0},
			},
		},
		{ // Middle of the month, no balance.
			ops:  Operations{NewOperation(Deposit, 1.0, now)},
			summ: Summaries{{Date: monthYear(now), Balance: 0, Change: 1.0}},
		},
		{ // Two balances --> Use the most recent.
			ops: Operations{
				NewOperation(Balance, 1.2, incMin),
				NewOperation(Balance, 2.2, now),
			},
			summ: Summaries{{Date: monthYear(now), Balance: 1.2, Change: 0}},
		},
	}
	for _, d := range data {
		got := d.ops.Summarize()
		if (len(got) != 0 && len(d.summ) != 0) && !reflect.DeepEqual(got, d.summ) {
			t.Errorf("got:%+v want:%+v", got, d.summ)
		}
	}
}
<commit_msg>Fixing tests.<commit_after>package isumm

import (
	"fmt"
	"reflect"
	"sort"
	"testing"
	"time"
)

var (
	now       = time.Now()
	nextMonth = now.Add(32 * 24 * time.Hour)
	incMin    = now.Add(10 * time.Minute)
	opType    = "1"
	value     = "12"
	date      = "2006-01-02"
)

func TestSort(t *testing.T) {
	o := Operations{NewOperation(Balance, 1, now), NewOperation(Balance, 1, incMin)}
	sort.Sort(o)
	if o[0].Date().Before(o[1].Date()) {
		t.Errorf("want %s before than %s", o[0], o[1])
	}
}

func TestNewOperationFromString(t *testing.T) {
	validOp, _ := NewOperationFromString(opType, value, date)
	testCases := []struct {
		desc, t, v, d string
		want          *Operation
	}{
		// Invalid cases.
		{desc: "invalid value - empty", t: opType, v: "", d: date},
		{desc: "invalid value - chars", t: opType, v: "acb", d: date},
		{desc: "invalid op - empty string", t: "", v: value, d: date},
		{desc: "invalid op - type does not exist", t: "94879138ddffg", v: value, d: date},
		{desc: "invalid date", t: opType, v: value, d: "31/31/31"},
		// Valid cases.
		{desc: "valid - contain spaces", t: opType, v: fmt.Sprintf(" %s ", value), d: date, want: &validOp},
		{desc: "valid - perfect", t: opType, v: value, d: date, want: &validOp},
	}
	for _, test := range testCases {
		got, err := NewOperationFromString(test.t, test.v, test.d)
		switch {
		case test.want == nil:
			if err == nil {
				t.Errorf("got:nil expected:err. Test:(%+v)", test)
			}
		default:
			if !reflect.DeepEqual(test.want, &got) {
				t.Errorf("got:(%+v) want:(%+v). Test:(%+v)", got, test.want, test)
			}
		}
	}
}

func TestSummarize(t *testing.T) {
	data := []struct {
		ops  Operations
		summ Summaries
	}{
		{ // Simple case: One entry containing only balance.
			ops:  Operations{NewOperation(Balance, 1.2, now)},
			summ: Summaries{
				{
					Date: monthYear(now),
					Balance: 1.2,
					Change: 0,
					SummaryOps: []SummaryOp{
						{Index: 0, Operation: NewOperation(Balance, 1.2, now)},
					},
				},
			},
		},
		{ // Empty case.
			ops:  Operations{},
			summ: Summaries{},
		},
		{
			ops: Operations{
				NewOperation(Balance, 1.2, incMin),
				NewOperation(Deposit, 1.0, now),
				NewOperation(Balance, 2.2, nextMonth),
			},
			summ: Summaries{
				{
					Date: monthYear(nextMonth),
					Balance: 2.2,
					Change: 0,
					SummaryOps: []SummaryOp{
						{Index: 0, Operation: NewOperation(Balance, 2.2, nextMonth)},
					},
				},
				{
					Date: monthYear(now),
					Balance: 1.2,
					Change: 1.0,
					SummaryOps: []SummaryOp{
						{Index: 1, Operation: NewOperation(Balance, 1.2, incMin)},
						{Index: 2, Operation: NewOperation(Deposit, 1.0, now)},
					},
				},
			},
		},
		{ // Middle of the month, no balance.
			ops:  Operations{NewOperation(Deposit, 1.0, now)},
			summ: Summaries{
				{
					Date: monthYear(now),
					Balance: 0,
					Change: 1.0,
					SummaryOps: []SummaryOp{
						{Index: 0, Operation: NewOperation(Deposit, 1.0, now)},
					},
				},
			},
		},
		{ // Two balances --> Use the most recent.
			ops: Operations{
				NewOperation(Balance, 1.2, incMin),
				NewOperation(Balance, 2.2, now),
			},
			summ: Summaries{
				{
					Date: monthYear(now),
					Balance: 1.2,
					Change: 0,
					SummaryOps: []SummaryOp{
						{Index: 0, Operation: NewOperation(Balance, 1.2, incMin)},
						{Index: 1, Operation: NewOperation(Balance, 2.2, now)},
					},
				},
			},
		},
	}
	for _, d := range data {
		got := d.ops.Summarize()
		if (len(got) != 0 && len(d.summ) != 0) && !reflect.DeepEqual(got, d.summ) {
			t.Errorf("got:%+v want:%+v", got, d.summ)
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>protocol/bc: make this benchmark test a little more focused<commit_after><|endoftext|>


<commit_before>package sync

import (
	"fmt"
	"log"
	"time"

	dm "github.com/svdberg/syncmysport-runkeeper/datamodel"
	rk "github.com/svdberg/syncmysport-runkeeper/runkeeper"
	stv "github.com/svdberg/syncmysport-runkeeper/strava"
)

type Syncer interface {
	Sync() (int, int, error)
}

/*
 * Returns the set of activities that are in RunKeeper, but not in Strava.
 * So if the set of Runkeeper activites is A, and the set of Strava activities is B,
 * this function calculates B\A.
 */
func CalculateRKDifference(rkActivities dm.ActivitySet, stvActivities dm.ActivitySet) *dm.ActivitySet {
	return rkActivities.ApproxSubtract(stvActivities)
}

type SyncTask struct {
	StravaToken       string `json:"stv_token"`
	RunkeeperToken    string `json:"rk_token"`
	LastSeenTimestamp int    `json:"last_seen_ts"`
	Uid               int64  `json:"id"`
	Environment       string `json:"environment"`
}

func CreateSyncTask(rkToken string, stvToken string, lastSeenTS int, environment string) *SyncTask {
	return &SyncTask{stvToken, rkToken, lastSeenTS, -1, environment}
}

/*
 * return the Total difference and the number of Activites created
 */
func (st SyncTask) Sync(stvClient stv.StravaClientInt, rkClient rk.RunkeeperCientInt) (int, int, error) {
	//get activities from strava
	//normalize time to the start of the day, because Runkeeper only supports days as offset, not timestamps
	tsOfStartOfDay := calculateTsAtStartOfDay(st.LastSeenTimestamp)
	activities, err := stvClient.GetSTVActivitiesSince(tsOfStartOfDay)
	if err != nil {
		log.Print("Error retrieving Strava activitites since %s, aborting this run", st.LastSeenTimestamp)
		return 0, 0, err
	}
	stvDetailedActivities := dm.NewActivitySet()
	for _, actSummary := range activities {
		//get Detailed Actv
		detailedAct, _ := stvClient.GetSTVDetailedActivity(actSummary.Id)
		//get associated Streams
		timeStream, err := stvClient.GetSTVActivityStream(actSummary.Id, "Time")
		if err != nil {
			log.Fatal("Error while retrieving time series from Strava: %s", err)
		}
		locStream, _ := stvClient.GetSTVActivityStream(actSummary.Id, "GPS")
		hrStream, _ := stvClient.GetSTVActivityStream(actSummary.Id, "Heartrate")

		stvDetailedActivities.Add(*stv.ConvertToActivity(detailedAct, timeStream, locStream, hrStream))
	}
	log.Printf("Got %d items from Strava", stvDetailedActivities.NumElements())
	for i := 0; i < stvDetailedActivities.NumElements(); i++ {
		log.Printf("Strava Activity: %s", stvDetailedActivities.Get(i))
	}

	//get activities from runkeeper
	rkActivitiesOverview, err := rkClient.GetRKActivitiesSince(st.LastSeenTimestamp)
	rkDetailActivities := rkClient.EnrichRKActivities(rkActivitiesOverview)
	//log.Printf("rk detail activities: %s", rkDetailActivities)

	rkActivities := dm.NewActivitySet()
	for _, item := range rkDetailActivities {
		rkActivities.Add(*rk.ConvertToActivity(&item))
	}
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("Got %d items from RunKeeper", rkActivities.NumElements())
	for i := 0; i < rkActivities.NumElements(); i++ {
		log.Printf("Runkeeper Activity: %s", rkActivities.Get(i))
	}

	//caclulate difference
	itemsToSyncToRk := stvDetailedActivities.ApproxSubtract(rkActivities)
	log.Printf("Difference between Runkeeper and Strava is %d items", itemsToSyncToRk.NumElements())

	//write to runkeeper
	totalItemsCreated := 0
	for i := 0; i < itemsToSyncToRk.NumElements(); i++ {
		log.Printf("Now storing item %s to RunKeeper", itemsToSyncToRk.Get(i))
		var (
			uri string
			err error
		)
		if st.Environment == "Prod" {
			uri, err = rkClient.PostActivity(rk.ConvertToRkActivity(itemsToSyncToRk.Get(i)))
		} else {
			log.Print("Assuming DEBUG/TEST mode, not actually writing to Runkeeper")
			err = nil
			uri = "fake_uri"
		}

		if err != nil {
			log.Fatal("Something failed during the write to Runkeeper: %s", err)
		}
		if uri != "" {
			log.Printf("URI of activity: %s", uri)
			totalItemsCreated++
		}
	}
	return itemsToSyncToRk.NumElements(), totalItemsCreated, nil
}

func calculateTsAtStartOfDay(timestamp int) int {
	timeAtTimestamp := time.Unix(int64(timestamp), 0).UTC()
	year, month, day := timeAtTimestamp.Date()
	//	Mon Jan 2 15:04:05 -0700 MST 2006
	ts, err := time.Parse("2006-1-2 15:04:05 MST", fmt.Sprintf("%d-%d-%d 00:00:00 UTC", year, month, day))
	if err != nil {
		log.Printf("Error parsing time from y/m/d to ts of current day: %s", err)
	}
	ts = ts.Add(time.Duration(1) * time.Minute)
	return int(ts.Unix())
}
<commit_msg>Some further fixes to the error handling in the jobs<commit_after>package sync

import (
	"fmt"
	"log"
	"time"

	dm "github.com/svdberg/syncmysport-runkeeper/datamodel"
	rk "github.com/svdberg/syncmysport-runkeeper/runkeeper"
	stv "github.com/svdberg/syncmysport-runkeeper/strava"
)

type Syncer interface {
	Sync() (int, int, error)
}

/*
 * Returns the set of activities that are in RunKeeper, but not in Strava.
 * So if the set of Runkeeper activites is A, and the set of Strava activities is B,
 * this function calculates B\A.
 */
func CalculateRKDifference(rkActivities dm.ActivitySet, stvActivities dm.ActivitySet) *dm.ActivitySet {
	return rkActivities.ApproxSubtract(stvActivities)
}

type SyncTask struct {
	StravaToken       string `json:"stv_token"`
	RunkeeperToken    string `json:"rk_token"`
	LastSeenTimestamp int    `json:"last_seen_ts"`
	Uid               int64  `json:"id"`
	Environment       string `json:"environment"`
}

func CreateSyncTask(rkToken string, stvToken string, lastSeenTS int, environment string) *SyncTask {
	return &SyncTask{stvToken, rkToken, lastSeenTS, -1, environment}
}

/*
 * return the Total difference and the number of Activites created
 */
func (st SyncTask) Sync(stvClient stv.StravaClientInt, rkClient rk.RunkeeperCientInt) (int, int, error) {
	//get activities from strava
	//normalize time to the start of the day, because Runkeeper only supports days as offset, not timestamps
	tsOfStartOfDay := calculateTsAtStartOfDay(st.LastSeenTimestamp)
	activities, err := stvClient.GetSTVActivitiesSince(tsOfStartOfDay)
	if err != nil {
		log.Print("Error retrieving Strava activitites since %s, aborting this run", st.LastSeenTimestamp)
		return 0, 0, err
	}
	stvDetailedActivities := dm.NewActivitySet()
	for _, actSummary := range activities {
		//get Detailed Actv
		detailedAct, _ := stvClient.GetSTVDetailedActivity(actSummary.Id)
		//get associated Streams
		timeStream, err := stvClient.GetSTVActivityStream(actSummary.Id, "Time")
		if err != nil {
			log.Printf("Error while retrieving time series from Strava: %s", err)
			return 0, 0, err
		}
		locStream, _ := stvClient.GetSTVActivityStream(actSummary.Id, "GPS")
		hrStream, _ := stvClient.GetSTVActivityStream(actSummary.Id, "Heartrate")

		stvDetailedActivities.Add(*stv.ConvertToActivity(detailedAct, timeStream, locStream, hrStream))
	}
	log.Printf("Got %d items from Strava", stvDetailedActivities.NumElements())
	for i := 0; i < stvDetailedActivities.NumElements(); i++ {
		log.Printf("Strava Activity: %s", stvDetailedActivities.Get(i))
	}

	//get activities from runkeeper
	rkActivitiesOverview, err := rkClient.GetRKActivitiesSince(st.LastSeenTimestamp)
	rkDetailActivities := rkClient.EnrichRKActivities(rkActivitiesOverview)
	//log.Printf("rk detail activities: %s", rkDetailActivities)

	rkActivities := dm.NewActivitySet()
	for _, item := range rkDetailActivities {
		rkActivities.Add(*rk.ConvertToActivity(&item))
	}
	if err != nil {
		log.Printf("%s", err)
	}
	log.Printf("Got %d items from RunKeeper", rkActivities.NumElements())
	for i := 0; i < rkActivities.NumElements(); i++ {
		log.Printf("Runkeeper Activity: %s", rkActivities.Get(i))
	}

	//caclulate difference
	itemsToSyncToRk := stvDetailedActivities.ApproxSubtract(rkActivities)
	log.Printf("Difference between Runkeeper and Strava is %d items", itemsToSyncToRk.NumElements())

	//write to runkeeper
	totalItemsCreated := 0
	for i := 0; i < itemsToSyncToRk.NumElements(); i++ {
		log.Printf("Now storing item %s to RunKeeper", itemsToSyncToRk.Get(i))
		var (
			uri string
			err error
		)
		if st.Environment == "Prod" {
			uri, err = rkClient.PostActivity(rk.ConvertToRkActivity(itemsToSyncToRk.Get(i)))
		} else {
			log.Print("Assuming DEBUG/TEST mode, not actually writing to Runkeeper")
			err = nil
			uri = "fake_uri"
		}

		if err != nil {
			log.Printf("Something failed during the write to Runkeeper: %s", err)
			return itemsToSyncToRk.NumElements(), totalItemsCreated, err
		}
		if uri != "" {
			log.Printf("URI of activity: %s", uri)
			totalItemsCreated++
		}
	}
	return itemsToSyncToRk.NumElements(), totalItemsCreated, nil
}

func calculateTsAtStartOfDay(timestamp int) int {
	timeAtTimestamp := time.Unix(int64(timestamp), 0).UTC()
	year, month, day := timeAtTimestamp.Date()
	//	Mon Jan 2 15:04:05 -0700 MST 2006
	ts, err := time.Parse("2006-1-2 15:04:05 MST", fmt.Sprintf("%d-%d-%d 00:00:00 UTC", year, month, day))
	if err != nil {
		log.Printf("Error parsing time from y/m/d to ts of current day: %s", err)
	}
	ts = ts.Add(time.Duration(1) * time.Minute)
	return int(ts.Unix())
}
<|endoftext|>


<commit_before>package termite

import (
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// A in-memory cache of attributes.
//
// Invariants: for all entries, we have their parent directories too
type AttributeCache struct {
	mutex      sync.RWMutex
	attributes map[string]*FileAttr
	cond       *sync.Cond
	busy       map[string]bool
	getter     func(name string) *FileAttr
	statter    func(name string) *os.FileInfo
}

func NewAttributeCache(getter func(n string) *FileAttr,
	statter func(n string) *os.FileInfo) *AttributeCache {
	me := &AttributeCache{
		attributes: make(map[string]*FileAttr),
		busy:       map[string]bool{},
	}
	me.cond = sync.NewCond(&me.mutex)
	me.getter = getter
	me.statter = statter
	return me
}

var paranoia = false

func (me *AttributeCache) Verify() {
	if !paranoia {
		return
	}
	me.mutex.RLock()
	defer me.mutex.RUnlock()
	me.verify()
}

func (me *AttributeCache) verify() {
	if !paranoia {
		return
	}
	for k, v := range me.attributes {
		if k != "" && filepath.Clean(k) != k {
			log.Panicf("Unclean path %q", k)
		}
		if v.Path != k {
			log.Panicf("attributes mismatch %q %#v", k, v)
		}
		if _, ok := me.busy[k]; ok {
			log.Panicf("busy and attributes entry for %q", k)
		}
		if v.Deletion() {
			log.Panicf("Attribute cache may not contain deletions %q", k)
		}
		if v.IsDirectory() && v.NameModeMap == nil {
			log.Panicf("dir has no NameModeMap %q", k)
		}
		for childName, mode := range v.NameModeMap {
			if strings.Contains(childName, "\000") || strings.Contains(childName, "/") || len(childName) == 0 {
				log.Panicf("%q has illegal child name %q: %o", k, childName, mode)
			}
			if mode == 0 {
				log.Panicf("child has 0 mode: %q.%q", k, childName)
			}
		}
		dir, base := SplitPath(k)
		if base != k {
			parent := me.attributes[dir]
			if v.Deletion() && parent != nil && parent.NameModeMap[base] != 0 {
				log.Panicf("Parent %q has entry for deleted %q", dir, base)
			}
			if !v.Deletion() && parent == nil {
				log.Panicf("Missing parent for %q", k)
			}
			if !v.Deletion() && parent.NameModeMap[base] == 0 {
				log.Panicf("Parent %q has no entry for %q", dir, base)
			}
		}
	}
}

func (me *AttributeCache) Have(name string) bool {
	me.mutex.RLock()
	defer me.mutex.RUnlock()
	_, ok := me.attributes[name]
	return ok
}

func (me *AttributeCache) Get(name string) (rep *FileAttr) {
	return me.get(name, false)
}

func (me *AttributeCache) GetDir(name string) (rep *FileAttr) {
	return me.get(name, true)
}

func (me *AttributeCache) localGet(name string, withdir bool) (rep *FileAttr) {
	me.mutex.RLock()
	defer me.mutex.RUnlock()

	rep, ok := me.attributes[name]
	if ok {
		return rep.Copy(withdir)
	}

	if name != "" {
		dir, base := SplitPath(name)
		dirAttr := me.attributes[dir]
		if dirAttr != nil && dirAttr.NameModeMap != nil && dirAttr.NameModeMap[base] == 0 {
			return &FileAttr{Path: name}
		}
	}
	return nil
}

func (me *AttributeCache) get(name string, withdir bool) (rep *FileAttr) {
	rep = me.localGet(name, withdir)
	if rep != nil {
		return rep
	}

	if name != "" {
		dir, base := SplitPath(name)
		dirAttr := me.get(dir, true)
		if dirAttr.NameModeMap != nil && dirAttr.NameModeMap[base] == 0 {
			return &FileAttr{Path: name}
		}
	}

	defer me.Verify()
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for me.busy[name] && me.attributes[name] == nil {
		me.cond.Wait()
	}
	rep, ok := me.attributes[name]
	if ok {
		return rep
	}
	me.busy[name] = true
	me.mutex.Unlock()

	rep = me.getter(name)

	me.mutex.Lock()
	if rep == nil {
		// This is an error, but what can we do?
		return &FileAttr{Path: name}
	}
	rep.Path = name

	if !rep.Deletion() {
		me.attributes[name] = rep
	}
	me.cond.Broadcast()
	me.busy[name] = false, false
	return rep.Copy(withdir)
}

func (me *AttributeCache) Update(files []*FileAttr) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.update(files)
}

func (me *AttributeCache) update(files []*FileAttr) {
	defer me.verify()
	attributes := me.attributes
	for _, inF := range files {
		r := *inF
		if len(r.Path) > 0 && r.Path[0] == '/' {
			panic("Leading slash.")
		}

		dir, basename := SplitPath(r.Path)
		if basename != "" {
			dirAttr := attributes[dir]
			if dirAttr == nil {
				log.Println("Discarding update: ", r)
				continue
			}
			if dirAttr.NameModeMap == nil {
				log.Panicf("parent dir has no NameModeMap: %q", dir)
			}
			if r.Deletion() {
				dirAttr.NameModeMap[basename] = 0, false
			} else {
				dirAttr.NameModeMap[basename] = FileMode(r.Mode &^ 07777)
			}
		}

		if r.Deletion() {
			attributes[r.Path] = nil, false
			continue
		}

		old := attributes[r.Path]
		if old == nil {
			old = &r
			attributes[r.Path] = old
		}
		old.Merge(r)
	}
}

func (me *AttributeCache) Refresh(prefix string) FileSet {
	me.mutex.Lock()
	defer me.mutex.Unlock()

	if prefix != "" && prefix[0] == '/' {
		panic("leading /")
	}

	updated := []*FileAttr{}
	for key, attr := range me.attributes {
		// TODO -should just do everything?
		if !strings.HasPrefix(key, prefix) {
			continue
		}

		fi := me.statter(key)
		if fi == nil && !attr.Deletion() {
			del := FileAttr{
				Path: key,
			}
			updated = append(updated, &del)
		}
		// TODO - does this handle symlinks corrrectly?
		if fi != nil && attr.FileInfo != nil && EncodeFileInfo(*attr.FileInfo) != EncodeFileInfo(*fi) {
			newEnt := me.getter(key)
			newEnt.Path = key
			updated = append(updated, newEnt)
		}
	}
	fs := FileSet{updated}
	fs.Sort()

	me.update(fs.Files)
	return fs
}

func (me *AttributeCache) Copy() FileSet {
	me.mutex.RLock()
	defer me.mutex.RUnlock()

	dump := []*FileAttr{}
	for _, attr := range me.attributes {
		dump = append(dump, attr.Copy(true))
	}

	fs := FileSet{dump}
	fs.Sort()
	return fs
}
<commit_msg>AttributeCache: nits:<commit_after>package termite

import (
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// A in-memory cache of attributes.
//
// Invariants: for all entries, we have their parent directories too
type AttributeCache struct {
	mutex      sync.RWMutex
	attributes map[string]*FileAttr
	cond       *sync.Cond
	busy       map[string]bool
	getter     func(name string) *FileAttr
	statter    func(name string) *os.FileInfo
}

func NewAttributeCache(getter func(n string) *FileAttr,
	statter func(n string) *os.FileInfo) *AttributeCache {
	me := &AttributeCache{
		attributes: make(map[string]*FileAttr),
		busy:       map[string]bool{},
	}
	me.cond = sync.NewCond(&me.mutex)
	me.getter = getter
	me.statter = statter

	return me
}

var paranoia = false

func (me *AttributeCache) Verify() {
	if !paranoia {
		return
	}
	me.mutex.RLock()
	defer me.mutex.RUnlock()
	me.verify()
}

func (me *AttributeCache) verify() {
	if !paranoia {
		return
	}
	for k, v := range me.attributes {
		if k != "" && filepath.Clean(k) != k {
			log.Panicf("Unclean path %q", k)
		}
		if v.Path != k {
			log.Panicf("attributes mismatch %q %#v", k, v)
		}
		if _, ok := me.busy[k]; ok {
			log.Panicf("busy and attributes entry for %q", k)
		}
		if v.Deletion() {
			log.Panicf("Attribute cache may not contain deletions %q", k)
		}
		if v.IsDirectory() && v.NameModeMap == nil {
			log.Panicf("dir has no NameModeMap %q", k)
		}
		for childName, mode := range v.NameModeMap {
			if strings.Contains(childName, "\000") || strings.Contains(childName, "/") || len(childName) == 0 {
				log.Panicf("%q has illegal child name %q: %o", k, childName, mode)
			}
			if mode == 0 {
				log.Panicf("child has 0 mode: %q.%q", k, childName)
			}
		}
		dir, base := SplitPath(k)
		if base != k {
			parent := me.attributes[dir]
			if v.Deletion() && parent != nil && parent.NameModeMap[base] != 0 {
				log.Panicf("Parent %q has entry for deleted %q", dir, base)
			}
			if !v.Deletion() && parent == nil {
				log.Panicf("Missing parent for %q", k)
			}
			if !v.Deletion() && parent.NameModeMap[base] == 0 {
				log.Panicf("Parent %q has no entry for %q", dir, base)
			}
		}
	}
}

func (me *AttributeCache) Have(name string) bool {
	me.mutex.RLock()
	defer me.mutex.RUnlock()
	_, ok := me.attributes[name]
	return ok
}

func (me *AttributeCache) Get(name string) (rep *FileAttr) {
	return me.get(name, false)
}

func (me *AttributeCache) GetDir(name string) (rep *FileAttr) {
	return me.get(name, true)
}

func (me *AttributeCache) localGet(name string, withdir bool) (rep *FileAttr) {
	me.mutex.RLock()
	defer me.mutex.RUnlock()

	rep, ok := me.attributes[name]
	if ok {
		return rep.Copy(withdir)
	}

	if name != "" {
		dir, base := SplitPath(name)
		dirAttr := me.attributes[dir]
		if dirAttr != nil && dirAttr.NameModeMap != nil && dirAttr.NameModeMap[base] == 0 {
			return &FileAttr{Path: name}
		}
	}
	return nil
}

func (me *AttributeCache) get(name string, withdir bool) (rep *FileAttr) {
	rep = me.localGet(name, withdir)

	if rep != nil {
		return rep
	}

	defer me.Verify()
	me.mutex.Lock()
	defer me.mutex.Unlock()
	return me.unsafeGet(name, withdir)
}

func (me *AttributeCache) unsafeGet(name string, withdir bool) (rep *FileAttr) {
	if name != "" {
		dir, base := SplitPath(name)
		dirAttr := me.unsafeGet(dir, true)
		if dirAttr.NameModeMap != nil && dirAttr.NameModeMap[base] == 0 {
			return &FileAttr{Path: name}
		}
	}

	for me.busy[name] && me.attributes[name] == nil {
		me.cond.Wait()
	}
	rep, ok := me.attributes[name]
	if ok {
		return rep
	}
	me.busy[name] = true
	me.mutex.Unlock()

	rep = me.getter(name)

	me.mutex.Lock()
	if rep == nil {
		// This is an error, but what can we do?
		return &FileAttr{Path: name}
	}
	rep.Path = name

	if !rep.Deletion() {
		me.attributes[name] = rep
	}
	me.cond.Broadcast()
	me.busy[name] = false, false
	return rep.Copy(withdir)
}

func (me *AttributeCache) Update(files []*FileAttr) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.update(files)
}

func (me *AttributeCache) update(files []*FileAttr) {
	defer me.verify()
	attributes := me.attributes
	for _, inF := range files {
		r := *inF
		if len(r.Path) > 0 && r.Path[0] == '/' {
			panic("Leading slash.")
		}

		dir, basename := SplitPath(r.Path)
		if basename != "" {
			dirAttr := attributes[dir]
			if dirAttr == nil {
				log.Println("Discarding update: ", r)
				continue
			}
			if dirAttr.NameModeMap == nil {
				log.Panicf("parent dir has no NameModeMap: %q", dir)
			}
			if r.Deletion() {
				dirAttr.NameModeMap[basename] = 0, false
			} else {
				dirAttr.NameModeMap[basename] = FileMode(r.Mode &^ 07777)
			}
		}

		if r.Deletion() {
			attributes[r.Path] = nil, false
			continue
		}

		old := attributes[r.Path]
		if old == nil {
			old = &r
			attributes[r.Path] = old
		}
		old.Merge(r)
		me.busy[r.Path] = false, false
	}
	me.cond.Broadcast()
}

func (me *AttributeCache) Refresh(prefix string) FileSet {
	me.mutex.Lock()
	defer me.mutex.Unlock()

	if prefix != "" && prefix[0] == '/' {
		panic("leading /")
	}

	updated := []*FileAttr{}
	for key, attr := range me.attributes {
		// TODO -should just do everything?
		if !strings.HasPrefix(key, prefix) {
			continue
		}

		fi := me.statter(key)
		if fi == nil && !attr.Deletion() {
			del := FileAttr{
				Path: key,
			}
			updated = append(updated, &del)
		}
		// TODO - does this handle symlinks corrrectly?
		if fi != nil && attr.FileInfo != nil && EncodeFileInfo(*attr.FileInfo) != EncodeFileInfo(*fi) {
			newEnt := me.getter(key)
			newEnt.Path = key
			updated = append(updated, newEnt)
		}
	}
	fs := FileSet{updated}
	fs.Sort()

	me.update(fs.Files)
	return fs
}

func (me *AttributeCache) Copy() FileSet {
	me.mutex.RLock()
	defer me.mutex.RUnlock()

	dump := []*FileAttr{}
	for _, attr := range me.attributes {
		dump = append(dump, attr.Copy(true))
	}

	fs := FileSet{dump}
	fs.Sort()
	return fs
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/ogier/pflag"
	"github.com/papertrail/remote_syslog2/papertrail"
	"github.com/papertrail/remote_syslog2/syslog"
	"github.com/papertrail/remote_syslog2/utils"
	"launchpad.net/goyaml"
)

const (
	MinimumRefreshInterval = (time.Duration(10) * time.Second)
	DefaultConfigFile      = "/etc/log_files.yml"
)

type LogFile struct {
	Path string
	Tag  string
}

type ConfigFile struct {
	Files       []LogFile
	Destination struct {
		Host     string `yaml:"host"`
		Port     int    `yaml:"port"`
		Protocol string `yaml:"protocol"`
	}
	Hostname string `yaml:"hostname"`
	//SetYAML is only called on pointers
	RefreshInterval *RefreshInterval `yaml:"new_file_check_interval"`
	ExcludeFiles    *RegexCollection `yaml:"exclude_files"`
	ExcludePatterns *RegexCollection `yaml:"exclude_patterns"`
}

type ConfigManager struct {
	Config    ConfigFile
	FlagFiles []LogFile
	Flags     struct {
		Hostname        string
		DestHost        string
		DestPort        int
		ConfigFile      string
		LogLevels       string
		DebugLogFile    string
		PidFile         string
		RefreshInterval RefreshInterval
		UseTCP          bool
		UseTLS          bool
		NoDaemonize     bool
		Severity        string
		Facility        string
		Poll            bool
	}
}

type RefreshInterval struct {
	Duration time.Duration
}

func (r *RefreshInterval) String() string {
	return fmt.Sprint(*r)
}

func (r *RefreshInterval) Set(value string) error {
	d, err := time.ParseDuration(value)

	if err != nil {
		return err
	}

	if d < MinimumRefreshInterval {
		return fmt.Errorf("refresh interval must be greater than %s", MinimumRefreshInterval)
	}
	r.Duration = d
	return nil
}

func (r *RefreshInterval) SetYAML(tag string, value interface{}) bool {
	err := r.Set(value.(string))
	if err != nil {
		return false
	}
	return true
}

type RegexCollection []*regexp.Regexp

func (r *RegexCollection) Set(value string) error {
	exp, err := regexp.Compile(value)
	if err != nil {
		return err
	}
	*r = append(*r, exp)
	return nil
}

func (r *RegexCollection) String() string {
	return fmt.Sprint(*r)
}

func (r *RegexCollection) SetYAML(tag string, value interface{}) bool {
	items, ok := value.([]interface{})

	if !ok {
		return false
	}

	for _, item := range items {
		s, ok := item.(string)

		if !ok {
			return false
		}

		err := r.Set(s)
		if err != nil {
			panic(fmt.Sprintf("Failed to compile regex expression \"%s\"", s))
		}
	}

	return true
}

func NewConfigManager() ConfigManager {
	cm := ConfigManager{}
	err := cm.Initialize()

	if err != nil {
		log.Criticalf("Failed to configure the application: %s", err)
		os.Exit(1)
	}

	return cm
}

func (cm *ConfigManager) Initialize() error {
	cm.Config.ExcludeFiles = &RegexCollection{}
	cm.Config.ExcludePatterns = &RegexCollection{}
	cm.parseFlags()

	err := cm.readConfig()
	if err != nil {
		return err
	}
	return nil
}

func (cm *ConfigManager) parseFlags() {
	pflag.StringVarP(&cm.Flags.ConfigFile, "configfile", "c", DefaultConfigFile, "Path to config")
	pflag.StringVarP(&cm.Flags.DestHost, "dest-host", "d", "", "Destination syslog hostname or IP")
	pflag.IntVarP(&cm.Flags.DestPort, "dest-port", "p", 0, "Destination syslog port")
	if utils.CanDaemonize {
		pflag.BoolVarP(&cm.Flags.NoDaemonize, "no-detach", "D", false, "Don't daemonize and detach from the terminal")
	} else {
		cm.Flags.NoDaemonize = true
	}
	pflag.StringVarP(&cm.Flags.Facility, "facility", "f", "user", "Facility")
	pflag.StringVar(&cm.Flags.Hostname, "hostname", "", "Local hostname to send from")
	pflag.StringVar(&cm.Flags.PidFile, "pid-file", "", "Location of the PID file")
	// --parse-syslog
	pflag.StringVarP(&cm.Flags.Severity, "severity", "s", "notice", "Severity")
	// --strip-color
	pflag.BoolVar(&cm.Flags.UseTCP, "tcp", false, "Connect via TCP (no TLS)")
	pflag.BoolVar(&cm.Flags.UseTLS, "tls", false, "Connect via TCP with TLS")
	pflag.BoolVar(&cm.Flags.Poll, "poll", false, "Detect changes by polling instead of inotify")
	pflag.Var(&cm.Flags.RefreshInterval, "new-file-check-interval", "How often to check for new files")
	_ = pflag.Bool("no-eventmachine-tail", false, "No action, provided for backwards compatibility")
	_ = pflag.Bool("eventmachine-tail", false, "No action, provided for backwards compatibility")
	pflag.StringVar(&cm.Flags.DebugLogFile, "debug-log-cfg", "", "the debug log file")
	pflag.StringVar(&cm.Flags.LogLevels, "log", "<root>=INFO", "\"logging configuration <root>=INFO;first=TRACE\"")
	pflag.Parse()
	for _, arg := range pflag.Args() {
		log := strings.Split(arg, ":")
		if len(log) == 2 {
			cm.FlagFiles = append(cm.FlagFiles, LogFile{Tag: log[0], Path: log[1]})
		} else {
			cm.FlagFiles = append(cm.FlagFiles, LogFile{Tag: "", Path: log[0]})
		}
	}
}

func (cm *ConfigManager) readConfig() error {
	log.Infof("Reading configuration file %s", cm.Flags.ConfigFile)
	err := cm.loadConfigFile()
	if err != nil {
		log.Errorf("%s", err)
		return err
	}
	return nil
}

func (cm *ConfigManager) loadConfigFile() error {
	file, err := ioutil.ReadFile(cm.Flags.ConfigFile)
	// don't error if the default config file isn't found
	if os.IsNotExist(err) && cm.Flags.ConfigFile == DefaultConfigFile {
		return nil
	}
	if err != nil {
		return fmt.Errorf("Could not read the config file: %s", err)
	}

	err = goyaml.Unmarshal(file, &cm.Config)
	if err != nil {
		return fmt.Errorf("Could not parse the config file: %s", err)
	}
	return nil
}

func (cm *ConfigManager) Daemonize() bool {
	return !cm.Flags.NoDaemonize
}

func (cm *ConfigManager) Hostname() string {
	switch {
	case cm.Flags.Hostname != "":
		return cm.Flags.Hostname
	case cm.Config.Hostname != "":
		return cm.Config.Hostname
	default:
		hostname, _ := os.Hostname()
		return hostname
	}
}

func (cm *ConfigManager) RootCAs() *x509.CertPool {
	if cm.DestProtocol() == "tls" && cm.DestHost() == "logs.papertrailapp.com" {
		return papertrail.RootCA()
	} else {
		return nil
	}
}

func (cm *ConfigManager) DestHost() string {
	switch {
	case cm.Flags.DestHost != "":
		return cm.Flags.DestHost
	case cm.Config.Destination.Host == "":
		log.Criticalf("No destination hostname specified")
		os.Exit(1)
	}
	return cm.Config.Destination.Host
}

func (cm ConfigManager) DestPort() int {
	switch {
	case cm.Flags.DestPort != 0:
		return cm.Flags.DestPort
	case cm.Config.Destination.Port != 0:
		return cm.Config.Destination.Port
	default:
		return 514
	}
}

func (cm *ConfigManager) DestProtocol() string {
	switch {
	case cm.Flags.UseTLS:
		return "tls"
	case cm.Flags.UseTCP:
		return "tcp"
	case cm.Config.Destination.Protocol != "":
		return cm.Config.Destination.Protocol
	default:
		return "udp"
	}
}

func (cm *ConfigManager) Severity() syslog.Priority {
	s, err := syslog.Severity(cm.Flags.Severity)
	if err != nil {
		log.Criticalf("%s is not a designated facility", cm.Flags.Severity)
		os.Exit(1)
	}
	return s
}

func (cm *ConfigManager) Facility() syslog.Priority {
	f, err := syslog.Facility(cm.Flags.Facility)
	if err != nil {
		log.Criticalf("%s is not a designated facility", cm.Flags.Facility)
		os.Exit(1)
	}
	return f
}

func (cm *ConfigManager) Poll() bool {
	return cm.Flags.Poll
}

func (cm *ConfigManager) Files() []LogFile {
	return append(cm.FlagFiles, cm.Config.Files...)
}

func (cm *ConfigManager) DebugLogFile() string {
	switch {
	case cm.Flags.DebugLogFile != "":
		return cm.Flags.DebugLogFile
	default:
		return "/dev/null"
	}
}

func (cm *ConfigManager) defaultPidFile() string {
	pidFiles := []string{
		"/var/run/remote_syslog.pid",
		os.Getenv("HOME") + "/run/remote_syslog.pid",
		os.Getenv("HOME") + "/tmp/remote_syslog.pid",
		os.Getenv("HOME") + "/remote_syslog.pid",
		os.TempDir() + "/remote_syslog.pid",
		os.Getenv("TMPDIR") + "/remote_syslog.pid",
	}
	for _, f := range pidFiles {
		dir := filepath.Dir(f)
		dirStat, err := os.Stat(dir)
		if err != nil || dirStat == nil || !dirStat.IsDir() {
			continue
		}
		fd, err := os.OpenFile(f, os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			continue
		}
		fd.Close()
		return f
	}
	return "/tmp/remote_syslog.pid"
}

func (cm *ConfigManager) PidFile() string {
	switch {
	case cm.Flags.PidFile != "":
		return cm.Flags.PidFile
	default:
		return cm.defaultPidFile()
	}
}

func (cm *ConfigManager) LogLevels() string {
	return cm.Flags.LogLevels
}

func (cm *ConfigManager) RefreshInterval() RefreshInterval {
	switch {
	case cm.Config.RefreshInterval != nil && cm.Flags.RefreshInterval.Duration != 0:
		return cm.Flags.RefreshInterval
	case cm.Config.RefreshInterval != nil:
		return *cm.Config.RefreshInterval
	case cm.Flags.RefreshInterval.Duration != 0:
		return cm.Flags.RefreshInterval
	}
	return RefreshInterval{Duration: MinimumRefreshInterval}
}

func (cm *ConfigManager) ExcludeFiles() []*regexp.Regexp {
	return *cm.Config.ExcludeFiles
}

func (cm *ConfigManager) ExcludePatterns() []*regexp.Regexp {
	return *cm.Config.ExcludePatterns
}
<commit_msg>Don't break the existing config file format.<commit_after>package main

import (
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/ogier/pflag"
	"github.com/papertrail/remote_syslog2/papertrail"
	"github.com/papertrail/remote_syslog2/syslog"
	"github.com/papertrail/remote_syslog2/utils"
	"launchpad.net/goyaml"
)

const (
	MinimumRefreshInterval = (time.Duration(10) * time.Second)
	DefaultConfigFile      = "/etc/log_files.yml"
)

type LogFile struct {
	Path string
	Tag  string
}

type ConfigFile struct {
	Files       []string
	Destination struct {
		Host     string `yaml:"host"`
		Port     int    `yaml:"port"`
		Protocol string `yaml:"protocol"`
	}
	Hostname string `yaml:"hostname"`
	//SetYAML is only called on pointers
	RefreshInterval *RefreshInterval `yaml:"new_file_check_interval"`
	ExcludeFiles    *RegexCollection `yaml:"exclude_files"`
	ExcludePatterns *RegexCollection `yaml:"exclude_patterns"`
}

type ConfigManager struct {
	Config    ConfigFile
	FlagFiles []LogFile
	Flags     struct {
		Hostname        string
		DestHost        string
		DestPort        int
		ConfigFile      string
		LogLevels       string
		DebugLogFile    string
		PidFile         string
		RefreshInterval RefreshInterval
		UseTCP          bool
		UseTLS          bool
		NoDaemonize     bool
		Severity        string
		Facility        string
		Poll            bool
	}
}

type RefreshInterval struct {
	Duration time.Duration
}

func (r *RefreshInterval) String() string {
	return fmt.Sprint(*r)
}

func (r *RefreshInterval) Set(value string) error {
	d, err := time.ParseDuration(value)

	if err != nil {
		return err
	}

	if d < MinimumRefreshInterval {
		return fmt.Errorf("refresh interval must be greater than %s", MinimumRefreshInterval)
	}
	r.Duration = d
	return nil
}

func (r *RefreshInterval) SetYAML(tag string, value interface{}) bool {
	err := r.Set(value.(string))
	if err != nil {
		return false
	}
	return true
}

type RegexCollection []*regexp.Regexp

func (r *RegexCollection) Set(value string) error {
	exp, err := regexp.Compile(value)
	if err != nil {
		return err
	}
	*r = append(*r, exp)
	return nil
}

func (r *RegexCollection) String() string {
	return fmt.Sprint(*r)
}

func (r *RegexCollection) SetYAML(tag string, value interface{}) bool {
	items, ok := value.([]interface{})

	if !ok {
		return false
	}

	for _, item := range items {
		s, ok := item.(string)

		if !ok {
			return false
		}

		err := r.Set(s)
		if err != nil {
			panic(fmt.Sprintf("Failed to compile regex expression \"%s\"", s))
		}
	}

	return true
}

func NewConfigManager() ConfigManager {
	cm := ConfigManager{}
	err := cm.Initialize()

	if err != nil {
		log.Criticalf("Failed to configure the application: %s", err)
		os.Exit(1)
	}

	return cm
}

func (cm *ConfigManager) Initialize() error {
	cm.Config.ExcludeFiles = &RegexCollection{}
	cm.Config.ExcludePatterns = &RegexCollection{}
	cm.parseFlags()

	err := cm.readConfig()
	if err != nil {
		return err
	}
	return nil
}

func (cm *ConfigManager) parseFlags() {
	pflag.StringVarP(&cm.Flags.ConfigFile, "configfile", "c", DefaultConfigFile, "Path to config")
	pflag.StringVarP(&cm.Flags.DestHost, "dest-host", "d", "", "Destination syslog hostname or IP")
	pflag.IntVarP(&cm.Flags.DestPort, "dest-port", "p", 0, "Destination syslog port")
	if utils.CanDaemonize {
		pflag.BoolVarP(&cm.Flags.NoDaemonize, "no-detach", "D", false, "Don't daemonize and detach from the terminal")
	} else {
		cm.Flags.NoDaemonize = true
	}
	pflag.StringVarP(&cm.Flags.Facility, "facility", "f", "user", "Facility")
	pflag.StringVar(&cm.Flags.Hostname, "hostname", "", "Local hostname to send from")
	pflag.StringVar(&cm.Flags.PidFile, "pid-file", "", "Location of the PID file")
	// --parse-syslog
	pflag.StringVarP(&cm.Flags.Severity, "severity", "s", "notice", "Severity")
	// --strip-color
	pflag.BoolVar(&cm.Flags.UseTCP, "tcp", false, "Connect via TCP (no TLS)")
	pflag.BoolVar(&cm.Flags.UseTLS, "tls", false, "Connect via TCP with TLS")
	pflag.BoolVar(&cm.Flags.Poll, "poll", false, "Detect changes by polling instead of inotify")
	pflag.Var(&cm.Flags.RefreshInterval, "new-file-check-interval", "How often to check for new files")
	_ = pflag.Bool("no-eventmachine-tail", false, "No action, provided for backwards compatibility")
	_ = pflag.Bool("eventmachine-tail", false, "No action, provided for backwards compatibility")
	pflag.StringVar(&cm.Flags.DebugLogFile, "debug-log-cfg", "", "the debug log file")
	pflag.StringVar(&cm.Flags.LogLevels, "log", "<root>=INFO", "\"logging configuration <root>=INFO;first=TRACE\"")
	pflag.Parse()
	for _, arg := range pflag.Args() {
		log := strings.Split(arg, ":")
		if len(log) == 2 {
			cm.FlagFiles = append(cm.FlagFiles, LogFile{Tag: log[0], Path: log[1]})
		} else {
			cm.FlagFiles = append(cm.FlagFiles, LogFile{Tag: "", Path: log[0]})
		}
	}
}

func (cm *ConfigManager) readConfig() error {
	log.Infof("Reading configuration file %s", cm.Flags.ConfigFile)
	err := cm.loadConfigFile()
	if err != nil {
		log.Errorf("%s", err)
		return err
	}
	return nil
}

func (cm *ConfigManager) loadConfigFile() error {
	file, err := ioutil.ReadFile(cm.Flags.ConfigFile)
	// don't error if the default config file isn't found
	if os.IsNotExist(err) && cm.Flags.ConfigFile == DefaultConfigFile {
		return nil
	}
	if err != nil {
		return fmt.Errorf("Could not read the config file: %s", err)
	}

	err = goyaml.Unmarshal(file, &cm.Config)
	if err != nil {
		return fmt.Errorf("Could not parse the config file: %s", err)
	}
	return nil
}

func (cm *ConfigManager) Daemonize() bool {
	return !cm.Flags.NoDaemonize
}

func (cm *ConfigManager) Hostname() string {
	switch {
	case cm.Flags.Hostname != "":
		return cm.Flags.Hostname
	case cm.Config.Hostname != "":
		return cm.Config.Hostname
	default:
		hostname, _ := os.Hostname()
		return hostname
	}
}

func (cm *ConfigManager) RootCAs() *x509.CertPool {
	if cm.DestProtocol() == "tls" && cm.DestHost() == "logs.papertrailapp.com" {
		return papertrail.RootCA()
	} else {
		return nil
	}
}

func (cm *ConfigManager) DestHost() string {
	switch {
	case cm.Flags.DestHost != "":
		return cm.Flags.DestHost
	case cm.Config.Destination.Host == "":
		log.Criticalf("No destination hostname specified")
		os.Exit(1)
	}
	return cm.Config.Destination.Host
}

func (cm ConfigManager) DestPort() int {
	switch {
	case cm.Flags.DestPort != 0:
		return cm.Flags.DestPort
	case cm.Config.Destination.Port != 0:
		return cm.Config.Destination.Port
	default:
		return 514
	}
}

func (cm *ConfigManager) DestProtocol() string {
	switch {
	case cm.Flags.UseTLS:
		return "tls"
	case cm.Flags.UseTCP:
		return "tcp"
	case cm.Config.Destination.Protocol != "":
		return cm.Config.Destination.Protocol
	default:
		return "udp"
	}
}

func (cm *ConfigManager) Severity() syslog.Priority {
	s, err := syslog.Severity(cm.Flags.Severity)
	if err != nil {
		log.Criticalf("%s is not a designated facility", cm.Flags.Severity)
		os.Exit(1)
	}
	return s
}

func (cm *ConfigManager) Facility() syslog.Priority {
	f, err := syslog.Facility(cm.Flags.Facility)
	if err != nil {
		log.Criticalf("%s is not a designated facility", cm.Flags.Facility)
		os.Exit(1)
	}
	return f
}

func (cm *ConfigManager) Poll() bool {
	return cm.Flags.Poll
}

func (cm *ConfigManager) Files() []LogFile {
	logFiles := cm.FlagFiles
	for _, file := range cm.Config.Files {
		log := strings.Split(file, ":")
		if len(log) == 2 {
			logFiles = append(logFiles, LogFile{Tag: log[0], Path: log[1]})
		} else {
			logFiles = append(logFiles, LogFile{Tag: "", Path: log[0]})
		}
	}
	return logFiles
}

func (cm *ConfigManager) DebugLogFile() string {
	switch {
	case cm.Flags.DebugLogFile != "":
		return cm.Flags.DebugLogFile
	default:
		return "/dev/null"
	}
}

func (cm *ConfigManager) defaultPidFile() string {
	pidFiles := []string{
		"/var/run/remote_syslog.pid",
		os.Getenv("HOME") + "/run/remote_syslog.pid",
		os.Getenv("HOME") + "/tmp/remote_syslog.pid",
		os.Getenv("HOME") + "/remote_syslog.pid",
		os.TempDir() + "/remote_syslog.pid",
		os.Getenv("TMPDIR") + "/remote_syslog.pid",
	}
	for _, f := range pidFiles {
		dir := filepath.Dir(f)
		dirStat, err := os.Stat(dir)
		if err != nil || dirStat == nil || !dirStat.IsDir() {
			continue
		}
		fd, err := os.OpenFile(f, os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			continue
		}
		fd.Close()
		return f
	}
	return "/tmp/remote_syslog.pid"
}

func (cm *ConfigManager) PidFile() string {
	switch {
	case cm.Flags.PidFile != "":
		return cm.Flags.PidFile
	default:
		return cm.defaultPidFile()
	}
}

func (cm *ConfigManager) LogLevels() string {
	return cm.Flags.LogLevels
}

func (cm *ConfigManager) RefreshInterval() RefreshInterval {
	switch {
	case cm.Config.RefreshInterval != nil && cm.Flags.RefreshInterval.Duration != 0:
		return cm.Flags.RefreshInterval
	case cm.Config.RefreshInterval != nil:
		return *cm.Config.RefreshInterval
	case cm.Flags.RefreshInterval.Duration != 0:
		return cm.Flags.RefreshInterval
	}
	return RefreshInterval{Duration: MinimumRefreshInterval}
}

func (cm *ConfigManager) ExcludeFiles() []*regexp.Regexp {
	return *cm.Config.ExcludeFiles
}

func (cm *ConfigManager) ExcludePatterns() []*regexp.Regexp {
	return *cm.Config.ExcludePatterns
}
<|endoftext|>


<commit_before>// +build proto

// RPC-based remote cache. Similar to HTTP but likely higher performance.
package cache

import (
	"bytes"
	"core"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

	pb "cache/proto/rpc_cache"
)

const maxErrors = 5

type rpcCache struct {
	connection *grpc.ClientConn
	client     pb.RpcCacheClient
	Writeable  bool
	Connected  bool
	Connecting bool
	OSName     string
	numErrors  int32
	timeout    time.Duration
	startTime  time.Time
}

func (cache *rpcCache) Store(target *core.BuildTarget, key []byte) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s in RPC cache...", target.Label)
		artifacts := []*pb.Artifact{}
		for out := range cacheArtifacts(target) {
			artifacts2, err := cache.loadArtifacts(target, out)
			if err != nil {
				log.Warning("RPC cache failed to load artifact %s: %s", out, err)
				cache.error()
				return
			}
			artifacts = append(artifacts, artifacts2...)
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) StoreExtra(target *core.BuildTarget, key []byte, file string) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s : %s in RPC cache...", target.Label, file)
		artifacts, err := cache.loadArtifacts(target, file)
		if err != nil {
			log.Warning("RPC cache failed to load artifact %s: %s", file, err)
			cache.error()
			return
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) loadArtifacts(target *core.BuildTarget, file string) ([]*pb.Artifact, error) {
	artifacts := []*pb.Artifact{}
	outDir := target.OutDir()
	root := path.Join(outDir, file)
	err := filepath.Walk(root, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			content, err := ioutil.ReadFile(name)
			if err != nil {
				return err
			}
			artifacts = append(artifacts, &pb.Artifact{
				Package: target.Label.PackageName,
				Target:  target.Label.Name,
				File:    name[len(outDir)+1:],
				Body:    content,
			})
		}
		return nil
	})
	return artifacts, err
}

func (cache *rpcCache) sendArtifacts(target *core.BuildTarget, key []byte, artifacts []*pb.Artifact) {
	req := pb.StoreRequest{Artifacts: artifacts, Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	ctx, cancel := context.WithTimeout(context.Background(), cache.timeout)
	defer cancel()
	resp, err := cache.client.Store(ctx, &req)
	if err != nil {
		log.Warning("Error communicating with RPC cache server: %s", err)
		cache.error()
	} else if !resp.Success {
		log.Warning("Failed to store artifacts in RPC cache for %s", target.Label)
	}
}

func (cache *rpcCache) Retrieve(target *core.BuildTarget, key []byte) bool {
	if !cache.isConnected() {
		return false
	}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	for out := range cacheArtifacts(target) {
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: out}
		req.Artifacts = append(req.Artifacts, &artifact)
	}
	// We can't tell from here if retrieval has been successful for a target with no outputs.
	// This is kind of weird but not actually disallowed, and we already have a test case for it,
	// so might as well try to get it right here.
	if len(req.Artifacts) == 0 {
		return false
	}
	return cache.retrieveArtifacts(target, &req, true)
}

func (cache *rpcCache) RetrieveExtra(target *core.BuildTarget, key []byte, file string) bool {
	if !cache.isConnected() {
		return false
	}
	artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: file}
	artifacts := []*pb.Artifact{&artifact}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH, Artifacts: artifacts}
	return cache.retrieveArtifacts(target, &req, false)
}

func (cache *rpcCache) retrieveArtifacts(target *core.BuildTarget, req *pb.RetrieveRequest, remove bool) bool {
	ctx, cancel := context.WithTimeout(context.Background(), cache.timeout)
	defer cancel()
	response, err := cache.client.Retrieve(ctx, req)
	if err != nil {
		log.Warning("Failed to retrieve artifacts for %s", target.Label)
		cache.error()
		return false
	} else if !response.Success {
		// Quiet, this is almost certainly just a 'not found'
		log.Debug("Couldn't retrieve artifacts for %s [key %s] from RPC cache", target.Label, base64.RawURLEncoding.EncodeToString(req.Hash))
		return false
	}
	// Remove any existing outputs first; this is important for cases where the output is a
	// directory, because we get back individual artifacts, and we need to make sure that
	// only the retrieved artifacts are present in the output.
	if remove {
		for _, out := range target.Outputs() {
			out := path.Join(target.OutDir(), out)
			if err := os.RemoveAll(out); err != nil {
				log.Error("Failed to remove artifact %s: %s", out, err)
				return false
			}
		}
	}
	for _, artifact := range response.Artifacts {
		if !cache.writeFile(target, artifact.File, artifact.Body) {
			return false
		}
	}
	// Sanity check: if we don't get anything back, assume it probably wasn't really a success.
	return len(response.Artifacts) > 0
}

func (cache *rpcCache) writeFile(target *core.BuildTarget, file string, body []byte) bool {
	out := path.Join(target.OutDir(), file)
	if err := os.MkdirAll(path.Dir(out), core.DirPermissions); err != nil {
		log.Warning("Failed to create directory for artifacts: %s", err)
		return false
	}
	if err := core.WriteFile(bytes.NewReader(body), out, fileMode(target)); err != nil {
		log.Warning("RPC cache failed to write file %s", err)
		return false
	}
	log.Debug("Retrieved %s - %s from RPC cache", target.Label, file)
	return true
}

func (cache *rpcCache) Clean(target *core.BuildTarget) {
	if cache.isConnected() && cache.Writeable {
		req := pb.DeleteRequest{Os: runtime.GOOS, Arch: runtime.GOARCH}
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name}
		req.Artifacts = []*pb.Artifact{&artifact}
		response, err := cache.client.Delete(context.Background(), &req)
		if err != nil || !response.Success {
			log.Errorf("Failed to remove %s from RPC cache", target.Label)
		}
	}
}

func (cache *rpcCache) connect(config *core.Configuration) {
	// Change grpc to log using our implementation
	grpclog.SetLogger(&grpcLogMabob{})
	log.Info("Connecting to RPC cache at %s", config.Cache.RpcUrl)
	opts := []grpc.DialOption{grpc.WithTimeout(cache.timeout)}
	if config.Cache.RpcPublicKey != "" || config.Cache.RpcCACert != "" || config.Cache.RpcSecure {
		auth, err := loadAuth(config.Cache.RpcCACert, config.Cache.RpcPublicKey, config.Cache.RpcPrivateKey)
		if err != nil {
			log.Warning("Failed to load RPC cache auth keys: %s", err)
			return
		}
		opts = append(opts, auth)
	} else {
		opts = append(opts, grpc.WithInsecure())
	}
	connection, err := grpc.Dial(config.Cache.RpcUrl, opts...)
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to connect to RPC cache: %s", err)
		return
	}
	// Note that we have to actually send it a message here to validate the connection;
	// Dial() only seems to return errors for superficial failures like syntactically invalid addresses,
	// it will return essentially immediately even if the server doesn't exist.
	healthclient := healthpb.NewHealthClient(connection)
	ctx, cancel := context.WithTimeout(context.Background(), cache.timeout)
	defer cancel()
	resp, err := healthclient.Check(ctx, &healthpb.HealthCheckRequest{Service: "plz-rpc-cache"})
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to contact RPC cache: %s", err)
	} else if resp.Status != healthpb.HealthCheckResponse_SERVING {
		cache.Connecting = false
		log.Warning("RPC cache says it is not serving (%d)", resp.Status)
	} else {
		cache.connection = connection
		cache.client = pb.NewRpcCacheClient(connection)
		cache.Connected = true
		cache.Connecting = false
		log.Info("RPC cache connected after %0.2fs", time.Since(cache.startTime).Seconds())
	}
}

// isConnected checks if the cache is connected. If it's still trying to connect it allows a
// very brief wait to give it a chance to come online.
func (cache *rpcCache) isConnected() bool {
	if cache.Connected {
		return true
	} else if !cache.Connecting {
		return false
	}
	ticker := time.NewTicker(10 * time.Millisecond)
	for i := 0; i < 5 && cache.Connecting; i++ {
		<-ticker.C
	}
	ticker.Stop()
	return cache.Connected
}

// error increments the error counter on the cache, and disables it if it gets too high.
// Note that after this it won't reconnect; we could try that but it probably isn't worth it
// (it's unlikely to restart in time if it's got a nontrivial set of artifacts to scan) and
// the user has probably been pestered by enough messages already.
func (cache *rpcCache) error() {
	if atomic.AddInt32(&cache.numErrors, 1) >= maxErrors && cache.Connected {
		log.Warning("Disabling RPC cache, looks like the connection has been lost")
		cache.Connected = false
	}
}

func newRpcCache(config *core.Configuration) (*rpcCache, error) {
	cache := &rpcCache{
		Writeable:  config.Cache.RpcWriteable,
		Connecting: true,
		timeout:    time.Duration(config.Cache.RpcTimeout) * time.Second,
		startTime:  time.Now(),
	}
	go cache.connect(config)
	return cache, nil
}

// grpcLogMabob is an implementation of grpc's logging interface using our backend.
type grpcLogMabob struct{}

func (g *grpcLogMabob) Fatal(args ...interface{})                 { log.Fatal(args...) }
func (g *grpcLogMabob) Fatalf(format string, args ...interface{}) { log.Fatalf(format, args...) }
func (g *grpcLogMabob) Fatalln(args ...interface{})               { log.Fatal(args...) }
func (g *grpcLogMabob) Print(args ...interface{})                 { log.Warning("%s", args) }
func (g *grpcLogMabob) Printf(format string, args ...interface{}) { log.Warning(format, args...) }
func (g *grpcLogMabob) Println(args ...interface{})               { log.Warning("%s", args) }

// loadAuth loads authentication credentials from a given pair of public / private key files.
func loadAuth(caCert, publicKey, privateKey string) (grpc.DialOption, error) {
	config := tls.Config{}
	if publicKey != "" {
		log.Debug("Loading client certificate from %s, key %s", publicKey, privateKey)
		cert, err := tls.LoadX509KeyPair(publicKey, privateKey)
		if err != nil {
			return nil, err
		}
		config.Certificates = []tls.Certificate{cert}
	}
	if caCert != "" {
		log.Debug("Reading CA cert file from %s", caCert)
		cert, err := ioutil.ReadFile(caCert)
		if err != nil {
			return nil, err
		}
		config.RootCAs = x509.NewCertPool()
		if !config.RootCAs.AppendCertsFromPEM(cert) {
			return nil, fmt.Errorf("Failed to add any PEM certificates from %s", caCert)
		}
	}
	return grpc.WithTransportCredentials(credentials.NewTLS(&config)), nil
}
<commit_msg>Remove unnecessary multiplication<commit_after>// +build proto

// RPC-based remote cache. Similar to HTTP but likely higher performance.
package cache

import (
	"bytes"
	"core"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

	pb "cache/proto/rpc_cache"
)

const maxErrors = 5

type rpcCache struct {
	connection *grpc.ClientConn
	client     pb.RpcCacheClient
	Writeable  bool
	Connected  bool
	Connecting bool
	OSName     string
	numErrors  int32
	timeout    time.Duration
	startTime  time.Time
}

func (cache *rpcCache) Store(target *core.BuildTarget, key []byte) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s in RPC cache...", target.Label)
		artifacts := []*pb.Artifact{}
		for out := range cacheArtifacts(target) {
			artifacts2, err := cache.loadArtifacts(target, out)
			if err != nil {
				log.Warning("RPC cache failed to load artifact %s: %s", out, err)
				cache.error()
				return
			}
			artifacts = append(artifacts, artifacts2...)
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) StoreExtra(target *core.BuildTarget, key []byte, file string) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s : %s in RPC cache...", target.Label, file)
		artifacts, err := cache.loadArtifacts(target, file)
		if err != nil {
			log.Warning("RPC cache failed to load artifact %s: %s", file, err)
			cache.error()
			return
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) loadArtifacts(target *core.BuildTarget, file string) ([]*pb.Artifact, error) {
	artifacts := []*pb.Artifact{}
	outDir := target.OutDir()
	root := path.Join(outDir, file)
	err := filepath.Walk(root, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			content, err := ioutil.ReadFile(name)
			if err != nil {
				return err
			}
			artifacts = append(artifacts, &pb.Artifact{
				Package: target.Label.PackageName,
				Target:  target.Label.Name,
				File:    name[len(outDir)+1:],
				Body:    content,
			})
		}
		return nil
	})
	return artifacts, err
}

func (cache *rpcCache) sendArtifacts(target *core.BuildTarget, key []byte, artifacts []*pb.Artifact) {
	req := pb.StoreRequest{Artifacts: artifacts, Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	ctx, cancel := context.WithTimeout(context.Background(), cache.timeout)
	defer cancel()
	resp, err := cache.client.Store(ctx, &req)
	if err != nil {
		log.Warning("Error communicating with RPC cache server: %s", err)
		cache.error()
	} else if !resp.Success {
		log.Warning("Failed to store artifacts in RPC cache for %s", target.Label)
	}
}

func (cache *rpcCache) Retrieve(target *core.BuildTarget, key []byte) bool {
	if !cache.isConnected() {
		return false
	}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	for out := range cacheArtifacts(target) {
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: out}
		req.Artifacts = append(req.Artifacts, &artifact)
	}
	// We can't tell from here if retrieval has been successful for a target with no outputs.
	// This is kind of weird but not actually disallowed, and we already have a test case for it,
	// so might as well try to get it right here.
	if len(req.Artifacts) == 0 {
		return false
	}
	return cache.retrieveArtifacts(target, &req, true)
}

func (cache *rpcCache) RetrieveExtra(target *core.BuildTarget, key []byte, file string) bool {
	if !cache.isConnected() {
		return false
	}
	artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: file}
	artifacts := []*pb.Artifact{&artifact}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH, Artifacts: artifacts}
	return cache.retrieveArtifacts(target, &req, false)
}

func (cache *rpcCache) retrieveArtifacts(target *core.BuildTarget, req *pb.RetrieveRequest, remove bool) bool {
	ctx, cancel := context.WithTimeout(context.Background(), cache.timeout)
	defer cancel()
	response, err := cache.client.Retrieve(ctx, req)
	if err != nil {
		log.Warning("Failed to retrieve artifacts for %s", target.Label)
		cache.error()
		return false
	} else if !response.Success {
		// Quiet, this is almost certainly just a 'not found'
		log.Debug("Couldn't retrieve artifacts for %s [key %s] from RPC cache", target.Label, base64.RawURLEncoding.EncodeToString(req.Hash))
		return false
	}
	// Remove any existing outputs first; this is important for cases where the output is a
	// directory, because we get back individual artifacts, and we need to make sure that
	// only the retrieved artifacts are present in the output.
	if remove {
		for _, out := range target.Outputs() {
			out := path.Join(target.OutDir(), out)
			if err := os.RemoveAll(out); err != nil {
				log.Error("Failed to remove artifact %s: %s", out, err)
				return false
			}
		}
	}
	for _, artifact := range response.Artifacts {
		if !cache.writeFile(target, artifact.File, artifact.Body) {
			return false
		}
	}
	// Sanity check: if we don't get anything back, assume it probably wasn't really a success.
	return len(response.Artifacts) > 0
}

func (cache *rpcCache) writeFile(target *core.BuildTarget, file string, body []byte) bool {
	out := path.Join(target.OutDir(), file)
	if err := os.MkdirAll(path.Dir(out), core.DirPermissions); err != nil {
		log.Warning("Failed to create directory for artifacts: %s", err)
		return false
	}
	if err := core.WriteFile(bytes.NewReader(body), out, fileMode(target)); err != nil {
		log.Warning("RPC cache failed to write file %s", err)
		return false
	}
	log.Debug("Retrieved %s - %s from RPC cache", target.Label, file)
	return true
}

func (cache *rpcCache) Clean(target *core.BuildTarget) {
	if cache.isConnected() && cache.Writeable {
		req := pb.DeleteRequest{Os: runtime.GOOS, Arch: runtime.GOARCH}
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name}
		req.Artifacts = []*pb.Artifact{&artifact}
		response, err := cache.client.Delete(context.Background(), &req)
		if err != nil || !response.Success {
			log.Errorf("Failed to remove %s from RPC cache", target.Label)
		}
	}
}

func (cache *rpcCache) connect(config *core.Configuration) {
	// Change grpc to log using our implementation
	grpclog.SetLogger(&grpcLogMabob{})
	log.Info("Connecting to RPC cache at %s", config.Cache.RpcUrl)
	opts := []grpc.DialOption{grpc.WithTimeout(cache.timeout)}
	if config.Cache.RpcPublicKey != "" || config.Cache.RpcCACert != "" || config.Cache.RpcSecure {
		auth, err := loadAuth(config.Cache.RpcCACert, config.Cache.RpcPublicKey, config.Cache.RpcPrivateKey)
		if err != nil {
			log.Warning("Failed to load RPC cache auth keys: %s", err)
			return
		}
		opts = append(opts, auth)
	} else {
		opts = append(opts, grpc.WithInsecure())
	}
	connection, err := grpc.Dial(config.Cache.RpcUrl, opts...)
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to connect to RPC cache: %s", err)
		return
	}
	// Note that we have to actually send it a message here to validate the connection;
	// Dial() only seems to return errors for superficial failures like syntactically invalid addresses,
	// it will return essentially immediately even if the server doesn't exist.
	healthclient := healthpb.NewHealthClient(connection)
	ctx, cancel := context.WithTimeout(context.Background(), cache.timeout)
	defer cancel()
	resp, err := healthclient.Check(ctx, &healthpb.HealthCheckRequest{Service: "plz-rpc-cache"})
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to contact RPC cache: %s", err)
	} else if resp.Status != healthpb.HealthCheckResponse_SERVING {
		cache.Connecting = false
		log.Warning("RPC cache says it is not serving (%d)", resp.Status)
	} else {
		cache.connection = connection
		cache.client = pb.NewRpcCacheClient(connection)
		cache.Connected = true
		cache.Connecting = false
		log.Info("RPC cache connected after %0.2fs", time.Since(cache.startTime).Seconds())
	}
}

// isConnected checks if the cache is connected. If it's still trying to connect it allows a
// very brief wait to give it a chance to come online.
func (cache *rpcCache) isConnected() bool {
	if cache.Connected {
		return true
	} else if !cache.Connecting {
		return false
	}
	ticker := time.NewTicker(10 * time.Millisecond)
	for i := 0; i < 5 && cache.Connecting; i++ {
		<-ticker.C
	}
	ticker.Stop()
	return cache.Connected
}

// error increments the error counter on the cache, and disables it if it gets too high.
// Note that after this it won't reconnect; we could try that but it probably isn't worth it
// (it's unlikely to restart in time if it's got a nontrivial set of artifacts to scan) and
// the user has probably been pestered by enough messages already.
func (cache *rpcCache) error() {
	if atomic.AddInt32(&cache.numErrors, 1) >= maxErrors && cache.Connected {
		log.Warning("Disabling RPC cache, looks like the connection has been lost")
		cache.Connected = false
	}
}

func newRpcCache(config *core.Configuration) (*rpcCache, error) {
	cache := &rpcCache{
		Writeable:  config.Cache.RpcWriteable,
		Connecting: true,
		timeout:    time.Duration(config.Cache.RpcTimeout),
		startTime:  time.Now(),
	}
	go cache.connect(config)
	return cache, nil
}

// grpcLogMabob is an implementation of grpc's logging interface using our backend.
type grpcLogMabob struct{}

func (g *grpcLogMabob) Fatal(args ...interface{})                 { log.Fatal(args...) }
func (g *grpcLogMabob) Fatalf(format string, args ...interface{}) { log.Fatalf(format, args...) }
func (g *grpcLogMabob) Fatalln(args ...interface{})               { log.Fatal(args...) }
func (g *grpcLogMabob) Print(args ...interface{})                 { log.Warning("%s", args) }
func (g *grpcLogMabob) Printf(format string, args ...interface{}) { log.Warning(format, args...) }
func (g *grpcLogMabob) Println(args ...interface{})               { log.Warning("%s", args) }

// loadAuth loads authentication credentials from a given pair of public / private key files.
func loadAuth(caCert, publicKey, privateKey string) (grpc.DialOption, error) {
	config := tls.Config{}
	if publicKey != "" {
		log.Debug("Loading client certificate from %s, key %s", publicKey, privateKey)
		cert, err := tls.LoadX509KeyPair(publicKey, privateKey)
		if err != nil {
			return nil, err
		}
		config.Certificates = []tls.Certificate{cert}
	}
	if caCert != "" {
		log.Debug("Reading CA cert file from %s", caCert)
		cert, err := ioutil.ReadFile(caCert)
		if err != nil {
			return nil, err
		}
		config.RootCAs = x509.NewCertPool()
		if !config.RootCAs.AppendCertsFromPEM(cert) {
			return nil, fmt.Errorf("Failed to add any PEM certificates from %s", caCert)
		}
	}
	return grpc.WithTransportCredentials(credentials.NewTLS(&config)), nil
}
<|endoftext|>


<commit_before>package shell

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle_map"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func init() {
	Commands = append(Commands, &commandVolumeFsck{})
}

type commandVolumeFsck struct {
	env *CommandEnv
}

func (c *commandVolumeFsck) Name() string {
	return "volume.fsck"
}

func (c *commandVolumeFsck) Help() string {
	return `check all volumes to find entries not used by the filer

	Important assumption!!!
		the system is all used by one filer.

	This command works this way:
	1. collect all file ids from all volumes, as set A
	2. collect all file ids from the filer, as set B
	3. find out the set A subtract B

`
}

func (c *commandVolumeFsck) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	c.env = commandEnv

	// collect all volume id locations
	volumeIdToServer, err := c.collectVolumeIds()
	if err != nil {
		return fmt.Errorf("failed to collect all volume locations: %v", err)
	}

	// create a temp folder
	tempFolder, err := ioutil.TempDir("", "sw_fsck")
	if err != nil {
		return fmt.Errorf("failed to create temp folder: %v", err)
	}
	fmt.Fprintf(writer, "working directory: %s\n", tempFolder)

	// collect each volume file ids
	for volumeId, vinfo := range volumeIdToServer {
		err = c.collectOneVolumeFileIds(tempFolder, volumeId, vinfo)
		if err != nil {
			return fmt.Errorf("failed to collect file ids from volume %d on %s: %v", volumeId, vinfo.server, err)
		}
	}

	// collect all filer file ids
	if err = c.collectFilerFileIds(tempFolder, volumeIdToServer); err != nil {
		return fmt.Errorf("failed to collect file ids from filer: %v", err)
	}

	// volume file ids substract filer file ids
	var totalOrphanChunkCount, totalOrphanDataSize uint64
	for volumeId, server := range volumeIdToServer {
		orphanChunkCount, orphanDataSize, checkErr := c.oneVolumeFileIdsSubtractFilerFileIds(tempFolder, volumeId)
		if checkErr != nil {
			return fmt.Errorf("failed to collect file ids from volume %d on %s: %v", volumeId, server, checkErr)
		}
		totalOrphanChunkCount += orphanChunkCount
		totalOrphanDataSize += orphanDataSize
	}

	if totalOrphanChunkCount > 0 {
		fmt.Fprintf(writer, "total %d orphan chunks, %d bytes\n", totalOrphanChunkCount, totalOrphanDataSize)
	} else {
		fmt.Fprintf(writer, "no orphan data\n")
	}

	os.RemoveAll(tempFolder)

	return nil
}

func (c *commandVolumeFsck) collectOneVolumeFileIds(tempFolder string, volumeId uint32, vinfo VInfo) error {

	return operation.WithVolumeServerClient(vinfo.server, c.env.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {

		copyFileClient, err := volumeServerClient.CopyFile(context.Background(), &volume_server_pb.CopyFileRequest{
			VolumeId:                 volumeId,
			Ext:                      ".idx",
			CompactionRevision:       math.MaxUint32,
			StopOffset:               math.MaxInt64,
			Collection:               vinfo.collection,
			IsEcVolume:               vinfo.isEcVolume,
			IgnoreSourceFileNotFound: false,
		})
		if err != nil {
			return fmt.Errorf("failed to start copying volume %d.idx: %v", volumeId, err)
		}

		err = writeToFile(copyFileClient, getVolumeFileIdFile(tempFolder, volumeId))
		if err != nil {
			return fmt.Errorf("failed to copy %s.idx from %s: %v", volumeId, vinfo.server, err)
		}

		return nil

	})

}

func (c *commandVolumeFsck) collectFilerFileIds(tempFolder string, volumeIdToServer map[uint32]VInfo) error {

	files := make(map[uint32]*os.File)
	for vid := range volumeIdToServer {
		dst, openErr := os.OpenFile(getFilerFileIdFile(tempFolder, vid), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if openErr != nil {
			return fmt.Errorf("failed to create file %s: %v", getFilerFileIdFile(tempFolder, vid), openErr)
		}
		files[vid] = dst
	}
	defer func() {
		for _, f := range files {
			f.Close()
		}
	}()

	type Item struct {
		vid     uint32
		fileKey uint64
	}
	return doTraverseBfsAndSaving(c.env, nil, "/", false, func(outputChan chan interface{}) {
		buffer := make([]byte, 8)
		for item := range outputChan {
			i := item.(*Item)
			util.Uint64toBytes(buffer, i.fileKey)
			files[i.vid].Write(buffer)
		}
	}, func(entry *filer_pb.FullEntry, outputChan chan interface{}) (err error) {
		for _, chunk := range entry.Entry.Chunks {
			outputChan <- &Item{
				vid:     chunk.Fid.VolumeId,
				fileKey: chunk.Fid.FileKey,
			}
		}
		return nil
	})
}

func (c *commandVolumeFsck) oneVolumeFileIdsSubtractFilerFileIds(tempFolder string, volumeId uint32) (orphanChunkCount, orphanDataSize uint64, err error) {

	db := needle_map.NewMemDb()
	defer db.Close()

	if err = db.LoadFromIdx(getVolumeFileIdFile(tempFolder, volumeId)); err != nil {
		return
	}

	filerFileIdsData, err := ioutil.ReadFile(getFilerFileIdFile(tempFolder, volumeId))
	if err != nil {
		return
	}

	dataLen := len(filerFileIdsData)
	if dataLen%8 != 0 {
		return 0, 0, fmt.Errorf("filer data is corrupted")
	}

	for i := 0; i < len(filerFileIdsData); i += 8 {
		fileKey := util.BytesToUint64(filerFileIdsData[i : i+8])
		db.Delete(types.NeedleId(fileKey))
	}

	db.AscendingVisit(func(n needle_map.NeedleValue) error {
		fmt.Printf("%d,%x\n", volumeId, n.Key)
		orphanChunkCount++
		orphanDataSize += uint64(n.Size)
		return nil
	})

	return

}

type VInfo struct {
	server     string
	collection string
	isEcVolume bool
}

func (c *commandVolumeFsck) collectVolumeIds() (volumeIdToServer map[uint32]VInfo, err error) {

	volumeIdToServer = make(map[uint32]VInfo)
	var resp *master_pb.VolumeListResponse
	err = c.env.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return
	}

	eachDataNode(resp.TopologyInfo, func(dc string, rack RackId, t *master_pb.DataNodeInfo) {
		for _, vi := range t.VolumeInfos {
			volumeIdToServer[vi.Id] = VInfo{
				server:     t.Id,
				collection: vi.Collection,
				isEcVolume: false,
			}
		}
		for _, ecShardInfo := range t.EcShardInfos {
			volumeIdToServer[ecShardInfo.Id] = VInfo{
				server:     t.Id,
				collection: ecShardInfo.Collection,
				isEcVolume: true,
			}
		}
	})

	return
}

func getVolumeFileIdFile(tempFolder string, vid uint32) string {
	return filepath.Join(tempFolder, fmt.Sprintf("%d.idx", vid))
}

func getFilerFileIdFile(tempFolder string, vid uint32) string {
	return filepath.Join(tempFolder, fmt.Sprintf("%d.fid", vid))
}

func writeToFile(client volume_server_pb.VolumeServer_CopyFileClient, fileName string) error {
	flags := os.O_WRONLY | os.O_CREATE | os.O_TRUNC
	dst, err := os.OpenFile(fileName, flags, 0644)
	if err != nil {
		return nil
	}
	defer dst.Close()

	for {
		resp, receiveErr := client.Recv()
		if receiveErr == io.EOF {
			break
		}
		if receiveErr != nil {
			return fmt.Errorf("receiving %s: %v", fileName, receiveErr)
		}
		dst.Write(resp.FileContent)
	}
	return nil
}
<commit_msg>better output format<commit_after>package shell

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle_map"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func init() {
	Commands = append(Commands, &commandVolumeFsck{})
}

type commandVolumeFsck struct {
	env *CommandEnv
}

func (c *commandVolumeFsck) Name() string {
	return "volume.fsck"
}

func (c *commandVolumeFsck) Help() string {
	return `check all volumes to find entries not used by the filer

	Important assumption!!!
		the system is all used by one filer.

	This command works this way:
	1. collect all file ids from all volumes, as set A
	2. collect all file ids from the filer, as set B
	3. find out the set A subtract B

`
}

func (c *commandVolumeFsck) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	c.env = commandEnv

	// collect all volume id locations
	volumeIdToServer, err := c.collectVolumeIds()
	if err != nil {
		return fmt.Errorf("failed to collect all volume locations: %v", err)
	}

	// create a temp folder
	tempFolder, err := ioutil.TempDir("", "sw_fsck")
	if err != nil {
		return fmt.Errorf("failed to create temp folder: %v", err)
	}
	// fmt.Fprintf(writer, "working directory: %s\n", tempFolder)

	// collect each volume file ids
	for volumeId, vinfo := range volumeIdToServer {
		err = c.collectOneVolumeFileIds(tempFolder, volumeId, vinfo)
		if err != nil {
			return fmt.Errorf("failed to collect file ids from volume %d on %s: %v", volumeId, vinfo.server, err)
		}
	}

	// collect all filer file ids
	if err = c.collectFilerFileIds(tempFolder, volumeIdToServer); err != nil {
		return fmt.Errorf("failed to collect file ids from filer: %v", err)
	}

	// volume file ids substract filer file ids
	var totalOrphanChunkCount, totalOrphanDataSize uint64
	for volumeId, server := range volumeIdToServer {
		orphanChunkCount, orphanDataSize, checkErr := c.oneVolumeFileIdsSubtractFilerFileIds(tempFolder, volumeId, writer)
		if checkErr != nil {
			return fmt.Errorf("failed to collect file ids from volume %d on %s: %v", volumeId, server, checkErr)
		}
		totalOrphanChunkCount += orphanChunkCount
		totalOrphanDataSize += orphanDataSize
	}

	if totalOrphanChunkCount > 0 {
		fmt.Fprintf(writer, "\ntotal\t%d orphan entries\t%d bytes not used by filer http://%s:%d/\n",
			totalOrphanChunkCount, totalOrphanDataSize, c.env.option.FilerHost, c.env.option.FilerPort)
		fmt.Fprintf(writer, "This could be normal if multiple filers or no filers are used.\n")
	} else {
		fmt.Fprintf(writer, "no orphan data\n")
	}

	os.RemoveAll(tempFolder)

	return nil
}

func (c *commandVolumeFsck) collectOneVolumeFileIds(tempFolder string, volumeId uint32, vinfo VInfo) error {

	return operation.WithVolumeServerClient(vinfo.server, c.env.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {

		copyFileClient, err := volumeServerClient.CopyFile(context.Background(), &volume_server_pb.CopyFileRequest{
			VolumeId:                 volumeId,
			Ext:                      ".idx",
			CompactionRevision:       math.MaxUint32,
			StopOffset:               math.MaxInt64,
			Collection:               vinfo.collection,
			IsEcVolume:               vinfo.isEcVolume,
			IgnoreSourceFileNotFound: false,
		})
		if err != nil {
			return fmt.Errorf("failed to start copying volume %d.idx: %v", volumeId, err)
		}

		err = writeToFile(copyFileClient, getVolumeFileIdFile(tempFolder, volumeId))
		if err != nil {
			return fmt.Errorf("failed to copy %s.idx from %s: %v", volumeId, vinfo.server, err)
		}

		return nil

	})

}

func (c *commandVolumeFsck) collectFilerFileIds(tempFolder string, volumeIdToServer map[uint32]VInfo) error {

	files := make(map[uint32]*os.File)
	for vid := range volumeIdToServer {
		dst, openErr := os.OpenFile(getFilerFileIdFile(tempFolder, vid), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if openErr != nil {
			return fmt.Errorf("failed to create file %s: %v", getFilerFileIdFile(tempFolder, vid), openErr)
		}
		files[vid] = dst
	}
	defer func() {
		for _, f := range files {
			f.Close()
		}
	}()

	type Item struct {
		vid     uint32
		fileKey uint64
	}
	return doTraverseBfsAndSaving(c.env, nil, "/", false, func(outputChan chan interface{}) {
		buffer := make([]byte, 8)
		for item := range outputChan {
			i := item.(*Item)
			util.Uint64toBytes(buffer, i.fileKey)
			files[i.vid].Write(buffer)
		}
	}, func(entry *filer_pb.FullEntry, outputChan chan interface{}) (err error) {
		for _, chunk := range entry.Entry.Chunks {
			outputChan <- &Item{
				vid:     chunk.Fid.VolumeId,
				fileKey: chunk.Fid.FileKey,
			}
		}
		return nil
	})
}

func (c *commandVolumeFsck) oneVolumeFileIdsSubtractFilerFileIds(tempFolder string, volumeId uint32, writer io.Writer) (orphanChunkCount, orphanDataSize uint64, err error) {

	db := needle_map.NewMemDb()
	defer db.Close()

	if err = db.LoadFromIdx(getVolumeFileIdFile(tempFolder, volumeId)); err != nil {
		return
	}

	filerFileIdsData, err := ioutil.ReadFile(getFilerFileIdFile(tempFolder, volumeId))
	if err != nil {
		return
	}

	dataLen := len(filerFileIdsData)
	if dataLen%8 != 0 {
		return 0, 0, fmt.Errorf("filer data is corrupted")
	}

	for i := 0; i < len(filerFileIdsData); i += 8 {
		fileKey := util.BytesToUint64(filerFileIdsData[i : i+8])
		db.Delete(types.NeedleId(fileKey))
	}

	db.AscendingVisit(func(n needle_map.NeedleValue) error {
		// fmt.Printf("%d,%x\n", volumeId, n.Key)
		orphanChunkCount++
		orphanDataSize += uint64(n.Size)
		return nil
	})

	if orphanChunkCount > 0 {
		fmt.Fprintf(writer, "volume %d\t%d orphan entries\t%d bytes\n", volumeId, orphanChunkCount, orphanDataSize)
	}

	return

}

type VInfo struct {
	server     string
	collection string
	isEcVolume bool
}

func (c *commandVolumeFsck) collectVolumeIds() (volumeIdToServer map[uint32]VInfo, err error) {

	volumeIdToServer = make(map[uint32]VInfo)
	var resp *master_pb.VolumeListResponse
	err = c.env.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return
	}

	eachDataNode(resp.TopologyInfo, func(dc string, rack RackId, t *master_pb.DataNodeInfo) {
		for _, vi := range t.VolumeInfos {
			volumeIdToServer[vi.Id] = VInfo{
				server:     t.Id,
				collection: vi.Collection,
				isEcVolume: false,
			}
		}
		for _, ecShardInfo := range t.EcShardInfos {
			volumeIdToServer[ecShardInfo.Id] = VInfo{
				server:     t.Id,
				collection: ecShardInfo.Collection,
				isEcVolume: true,
			}
		}
	})

	return
}

func getVolumeFileIdFile(tempFolder string, vid uint32) string {
	return filepath.Join(tempFolder, fmt.Sprintf("%d.idx", vid))
}

func getFilerFileIdFile(tempFolder string, vid uint32) string {
	return filepath.Join(tempFolder, fmt.Sprintf("%d.fid", vid))
}

func writeToFile(client volume_server_pb.VolumeServer_CopyFileClient, fileName string) error {
	flags := os.O_WRONLY | os.O_CREATE | os.O_TRUNC
	dst, err := os.OpenFile(fileName, flags, 0644)
	if err != nil {
		return nil
	}
	defer dst.Close()

	for {
		resp, receiveErr := client.Recv()
		if receiveErr == io.EOF {
			break
		}
		if receiveErr != nil {
			return fmt.Errorf("receiving %s: %v", fileName, receiveErr)
		}
		dst.Write(resp.FileContent)
	}
	return nil
}
<|endoftext|>


<commit_before>package resp

import (
	"bufio"
	"fmt"
	"io"
	"net"
)

// RedisConnection is the protocol reading and writing layer
type RedisConnection struct {
	Conn   net.Conn
	reader *bufio.Reader
}

func NewRedisConnection(conn net.Conn) *RedisConnection {
	return &RedisConnection{
		Conn:   conn,
		reader: bufio.NewReader(conn),
	}
}

func (rconn *RedisConnection) Close() error {
	return rconn.Conn.Close()
}

func (rconn *RedisConnection) Handle(handler RedisHandler) {
	if err := rconn.handle(handler); err != nil {
		fmt.Printf("redis handler error from %v: %v\n", rconn.Conn.RemoteAddr(), err)
	}
	if err := rconn.Close(); err != nil {
		fmt.Printf("error closing connection from %v: %v\n", rconn.Conn.RemoteAddr(), err)
	}
}

func (rconn *RedisConnection) handle(handler RedisHandler) error {
	if err := handler.HandleStart(rconn); err != nil {
		return err
	}

	for {
		err := rconn.Consume(handler)
		if err != nil {
			if err != io.EOF {
				rconn.WriteError(err)
			}
			break
		}
	}

	return handler.HandleEnd(rconn)
}

func (rconn *RedisConnection) Consume(handler RedisHandler) error {
	c, err := rconn.reader.ReadByte()
	if err != nil {
		return err
	}

	switch c {
	case '-':
		return rconn.consumeError(handler)

	case ':':
		return rconn.consumeInteger(handler)

	case '+':
		return rconn.consumeShortString(handler)

	case '$':
		return rconn.consumeBulkString(handler)

	case '*':
		return rconn.consumeArray(handler)

	default:
		return fmt.Errorf("unknown RESP type %#v", string(c))
	}

	return nil
}

func (rconn *RedisConnection) consumeError(handler RedisHandler) error {
	buf, err := rconn.scanLine()
	if err != nil {
		return err
	}
	return handler.HandleError(rconn, buf)
}

func (rconn *RedisConnection) consumeInteger(handler RedisHandler) error {
	n, err := rconn.readInteger()
	if err != nil {
		return err
	}
	return handler.HandleInteger(rconn, n)
}

func (rconn *RedisConnection) consumeShortString(handler RedisHandler) error {
	buf, err := rconn.scanLine()
	if err != nil {
		return err
	}
	return handler.HandleString(rconn, buf)
}

func (rconn *RedisConnection) consumeBulkString(handler RedisHandler) error {
	n, err := rconn.readInteger()
	if err != nil {
		return err
	}

	if n < 0 {
		return handler.HandleNull(rconn)
	}

	strReader := io.LimitReader(rconn.reader, int64(n))
	if err := handler.HandleBulkString(rconn, n, strReader); err != nil {
		return err
	}

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return err
	}
	if c != '\r' {
		return fmt.Errorf("missing CR")
	}

	c, err = rconn.reader.ReadByte()
	if err != nil {
		return err
	}
	if c != '\n' {
		return fmt.Errorf("missing LF after CR")
	}

	return nil
}

func (rconn *RedisConnection) consumeArray(handler RedisHandler) error {
	n, err := rconn.readInteger()
	if err != nil {
		return err
	}

	if n < 0 {
		return handler.HandleNull(rconn)
	}

	return handler.HandleArray(rconn, n)
}

func (rconn *RedisConnection) readUInteger() (uint, error) {
	n, err := rconn.scanNumbers(0)
	if err != nil {
		return n, err
	}

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return n, err
	}
	if c != '\n' {
		return n, fmt.Errorf("missing LF after CR")
	}

	return n, nil
}

func (rconn *RedisConnection) readInteger() (int, error) {
	n := 0

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return n, err
	}

	if c == '-' {
		nu, err := rconn.scanNumbers(0)
		n = -int(nu)
		if err != nil {
			return n, err
		}
	} else if c != '\r' {
		// if c > '9' {
		// 	return n, fmt.Errorf("unexpected byte %v while scanning integer, expected [0-9]", c)
		// }
		nu, err := rconn.scanNumbers(uint(c - '0'))
		n = int(nu)
		if err != nil {
			return n, err
		}
	}

	c, err = rconn.reader.ReadByte()
	if err != nil {
		return n, err
	}
	if c != '\n' {
		return n, fmt.Errorf("missing LF after CR")
	}

	return n, nil
}

func (rconn *RedisConnection) scanNumbers(n uint) (uint, error) {
	for {
		c, err := rconn.reader.ReadByte()
		if err != nil {
			return n, err
		}
		if c == '\r' {
			break
		}
		// if c > '9' {
		// 	return n, fmt.Errorf("unexpected byte %v while scanning integer, expected [0-9]", c)
		// }

		n = 10*n + uint(c-'0')
	}

	return n, nil
}

func (rconn *RedisConnection) scanLine() ([]byte, error) {
	buf, err := rconn.reader.ReadBytes('\r')
	if err != nil {
		return buf, err
	}

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return buf, err
	}
	if c != '\n' {
		return buf, fmt.Errorf("missing LF after CR")
	}

	return buf, nil
}

func (rconn *RedisConnection) WriteArrayHeader(num int) error {
	return rconn.writef("*%v\r\n", num)
}

func (rconn *RedisConnection) WriteInteger(num int) error {
	return rconn.writef(":%v\r\n", num)
}

func (rconn *RedisConnection) WriteNull() error {
	return rconn.write([]byte("$-1\r\n"))
}

func (rconn *RedisConnection) WriteNullArray() error {
	return rconn.write([]byte("*-1\r\n"))
}

func (rconn *RedisConnection) WriteBulkBytes(buf []byte) error {
	n := len(buf)
	if n == 0 {
		return rconn.write([]byte("$0\r\n\r\n"))
	}

	if err := rconn.writef("$%v\r\n", n); err != nil {
		return err
	}

	if _, err := rconn.Conn.Write(buf); err != nil {
		return err
	}

	return rconn.write([]byte("\r\n"))
}

func (rconn *RedisConnection) WriteBulkStringHeader(n int) error {
	return rconn.writef("$%v\r\n", n)
}

func (rconn *RedisConnection) WriteBulkStringFooter() error {
	return rconn.write([]byte("\r\n"))
}

func (rconn *RedisConnection) WriteBulkString(str string) error {
	n := len(str)
	if n == 0 {
		return rconn.write([]byte("$0\r\n\r\n"))
	}
	return rconn.writef("$%v\r\n%v\r\n", n, str)
}

func (rconn *RedisConnection) WriteSimpleString(str string) error {
	return rconn.writef("+%v\r\n", str)
}

func (rconn *RedisConnection) WriteSimpleBytes(b []byte) error {
	return rconn.writef("+%s\r\n", b)
}

func (rconn *RedisConnection) WriteError(err error) error {
	return rconn.writef("-ERR %v\r\n", err)
}

func (rconn *RedisConnection) WriteErrorBytes(b []byte) error {
	if _, err := rconn.Conn.Write([]byte("-")); err != nil {
		return err
	}
	if _, err := rconn.Conn.Write(b); err != nil {
		return err
	}
	if _, err := rconn.Conn.Write([]byte("\r\n")); err != nil {
		return err
	}
	return nil
}

func (rconn *RedisConnection) WriteErrorString(errType string, str string) error {
	return rconn.writef("-%v %v\r\n", errType, str)
}

func (rconn *RedisConnection) writef(format string, a ...interface{}) error {
	_, err := fmt.Fprintf(rconn.Conn, format, a...)
	return err
}

func (rconn *RedisConnection) write(buf []byte) error {
	if _, err := rconn.Conn.Write(buf); err != nil {
		return err
	}
	return nil
}
<commit_msg>RedisConnection: docs<commit_after>package resp

import (
	"bufio"
	"fmt"
	"io"
	"net"
)

// RedisConnection is the protocol reading and writing layer
type RedisConnection struct {
	Conn   net.Conn
	reader *bufio.Reader
}

// NewRedisConnection creates a redis connection around an existing net.Conn.
func NewRedisConnection(conn net.Conn) *RedisConnection {
	return &RedisConnection{
		Conn:   conn,
		reader: bufio.NewReader(conn),
	}
}

// Close closes the underlying connection.
func (rconn *RedisConnection) Close() error {
	return rconn.Conn.Close()
}

// Handle runs the passed handler until the connection ends or errors.
func (rconn *RedisConnection) Handle(handler RedisHandler) {
	if err := rconn.handle(handler); err != nil {
		fmt.Printf("redis handler error from %v: %v\n", rconn.Conn.RemoteAddr(), err)
	}
	if err := rconn.Close(); err != nil {
		fmt.Printf("error closing connection from %v: %v\n", rconn.Conn.RemoteAddr(), err)
	}
}

func (rconn *RedisConnection) handle(handler RedisHandler) error {
	if err := handler.HandleStart(rconn); err != nil {
		return err
	}

	for {
		err := rconn.Consume(handler)
		if err != nil {
			if err != io.EOF {
				rconn.WriteError(err)
			}
			break
		}
	}

	return handler.HandleEnd(rconn)
}

// Consume reads one element from the connection and passes it to the given handler.
func (rconn *RedisConnection) Consume(handler RedisHandler) error {
	c, err := rconn.reader.ReadByte()
	if err != nil {
		return err
	}

	switch c {
	case '-':
		return rconn.consumeError(handler)

	case ':':
		return rconn.consumeInteger(handler)

	case '+':
		return rconn.consumeShortString(handler)

	case '$':
		return rconn.consumeBulkString(handler)

	case '*':
		return rconn.consumeArray(handler)

	default:
		return fmt.Errorf("unknown RESP type %#v", string(c))
	}

	return nil
}

func (rconn *RedisConnection) consumeError(handler RedisHandler) error {
	buf, err := rconn.scanLine()
	if err != nil {
		return err
	}
	return handler.HandleError(rconn, buf)
}

func (rconn *RedisConnection) consumeInteger(handler RedisHandler) error {
	n, err := rconn.readInteger()
	if err != nil {
		return err
	}
	return handler.HandleInteger(rconn, n)
}

func (rconn *RedisConnection) consumeShortString(handler RedisHandler) error {
	buf, err := rconn.scanLine()
	if err != nil {
		return err
	}
	return handler.HandleString(rconn, buf)
}

func (rconn *RedisConnection) consumeBulkString(handler RedisHandler) error {
	n, err := rconn.readInteger()
	if err != nil {
		return err
	}

	if n < 0 {
		return handler.HandleNull(rconn)
	}

	strReader := io.LimitReader(rconn.reader, int64(n))
	if err := handler.HandleBulkString(rconn, n, strReader); err != nil {
		return err
	}

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return err
	}
	if c != '\r' {
		return fmt.Errorf("missing CR")
	}

	c, err = rconn.reader.ReadByte()
	if err != nil {
		return err
	}
	if c != '\n' {
		return fmt.Errorf("missing LF after CR")
	}

	return nil
}

func (rconn *RedisConnection) consumeArray(handler RedisHandler) error {
	n, err := rconn.readInteger()
	if err != nil {
		return err
	}

	if n < 0 {
		return handler.HandleNull(rconn)
	}

	return handler.HandleArray(rconn, n)
}

func (rconn *RedisConnection) readUInteger() (uint, error) {
	n, err := rconn.scanNumbers(0)
	if err != nil {
		return n, err
	}

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return n, err
	}
	if c != '\n' {
		return n, fmt.Errorf("missing LF after CR")
	}

	return n, nil
}

func (rconn *RedisConnection) readInteger() (int, error) {
	n := 0

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return n, err
	}

	if c == '-' {
		nu, err := rconn.scanNumbers(0)
		n = -int(nu)
		if err != nil {
			return n, err
		}
	} else if c != '\r' {
		// if c > '9' {
		// 	return n, fmt.Errorf("unexpected byte %v while scanning integer, expected [0-9]", c)
		// }
		nu, err := rconn.scanNumbers(uint(c - '0'))
		n = int(nu)
		if err != nil {
			return n, err
		}
	}

	c, err = rconn.reader.ReadByte()
	if err != nil {
		return n, err
	}
	if c != '\n' {
		return n, fmt.Errorf("missing LF after CR")
	}

	return n, nil
}

func (rconn *RedisConnection) scanNumbers(n uint) (uint, error) {
	for {
		c, err := rconn.reader.ReadByte()
		if err != nil {
			return n, err
		}
		if c == '\r' {
			break
		}
		// if c > '9' {
		// 	return n, fmt.Errorf("unexpected byte %v while scanning integer, expected [0-9]", c)
		// }

		n = 10*n + uint(c-'0')
	}

	return n, nil
}

func (rconn *RedisConnection) scanLine() ([]byte, error) {
	buf, err := rconn.reader.ReadBytes('\r')
	if err != nil {
		return buf, err
	}

	c, err := rconn.reader.ReadByte()
	if err != nil {
		return buf, err
	}
	if c != '\n' {
		return buf, fmt.Errorf("missing LF after CR")
	}

	return buf, nil
}

// WriteArrayHeader writes a "*N\r\n" array header.
func (rconn *RedisConnection) WriteArrayHeader(num int) error {
	return rconn.writef("*%v\r\n", num)
}

// WriteInteger writes a ":N\r\n" integer literal
func (rconn *RedisConnection) WriteInteger(num int) error {
	return rconn.writef(":%v\r\n", num)
}

// WriteNull writes a "$-1\r\n" null string
func (rconn *RedisConnection) WriteNull() error {
	return rconn.write([]byte("$-1\r\n"))
}

// WriteNullArray writes a "*-1\r\n" null array
func (rconn *RedisConnection) WriteNullArray() error {
	return rconn.write([]byte("*-1\r\n"))
}

// WriteBulkBytes writes "$N\r\n...\r\n" bulk string from a byte slice.
func (rconn *RedisConnection) WriteBulkBytes(buf []byte) error {
	n := len(buf)
	if n == 0 {
		return rconn.write([]byte("$0\r\n\r\n"))
	}

	if err := rconn.writef("$%v\r\n", n); err != nil {
		return err
	}

	if _, err := rconn.Conn.Write(buf); err != nil {
		return err
	}

	return rconn.write([]byte("\r\n"))
}

// WriteBulkStringHeader writes a "$N\r\n" bulk string header.
func (rconn *RedisConnection) WriteBulkStringHeader(n int) error {
	return rconn.writef("$%v\r\n", n)
}

// WriteBulkStringFooter writes a "\r\n" bulk string footer.
func (rconn *RedisConnection) WriteBulkStringFooter() error {
	return rconn.write([]byte("\r\n"))
}

// WriteBulkString writes a "$N\r\n...\r\n" bulk string.
func (rconn *RedisConnection) WriteBulkString(str string) error {
	n := len(str)
	if n == 0 {
		return rconn.write([]byte("$0\r\n\r\n"))
	}
	return rconn.writef("$%v\r\n%v\r\n", n, str)
}

// WriteSimpleString writes a "+...\r\n" simple string.
func (rconn *RedisConnection) WriteSimpleString(str string) error {
	return rconn.writef("+%v\r\n", str)
}

// WriteSimpleBytes writes a "+...\r\n" simple string froma byte slice.
func (rconn *RedisConnection) WriteSimpleBytes(b []byte) error {
	return rconn.writef("+%s\r\n", b)
}

// WriteError writes a "-ERR...\r\n" error.
func (rconn *RedisConnection) WriteError(err error) error {
	return rconn.writef("-ERR %v\r\n", err)
}

// WriteErrorBytes writes a "-...\r\n" error from a byte slice.
func (rconn *RedisConnection) WriteErrorBytes(b []byte) error {
	if _, err := rconn.Conn.Write([]byte("-")); err != nil {
		return err
	}
	if _, err := rconn.Conn.Write(b); err != nil {
		return err
	}
	if _, err := rconn.Conn.Write([]byte("\r\n")); err != nil {
		return err
	}
	return nil
}

// WriteErrorString writes a "-TYPE ...\r\n" error from a string type and body.
func (rconn *RedisConnection) WriteErrorString(errType, str string) error {
	return rconn.writef("-%v %v\r\n", errType, str)
}

func (rconn *RedisConnection) writef(format string, a ...interface{}) error {
	_, err := fmt.Fprintf(rconn.Conn, format, a...)
	return err
}

func (rconn *RedisConnection) write(buf []byte) error {
	if _, err := rconn.Conn.Write(buf); err != nil {
		return err
	}
	return nil
}
<|endoftext|>