add errors; add mutex; fix some old tests

src/go/mongolib/stats/stats.go

@@ -4,6 +4,7 @@ import (
 	"crypto/md5"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/montanaflynn/stats"
@@ -12,6 +13,25 @@ import (
 	"github.com/percona/percona-toolkit/src/go/mongolib/util"
 )
 
+type StatsError struct {
+	error
+}
+
+func (e *StatsError) Error() string {
+	if e == nil {
+		return "<nil>"
+	}
+
+	return fmt.Sprintf("stats error: %s", e.error)
+}
+
+func (e *StatsError) Parent() error {
+	return e.error
+}
+
+type StatsFingerprintError StatsError
+type StatsGetQueryFieldError StatsError
+
 // New creates new instance of stats with given fingerprinter
 func New(fingerprinter fingerprinter.Fingerprinter) *Stats {
 	s := &Stats{
@@ -29,10 +49,14 @@ type Stats struct {
 
 	// internal
 	queryInfoAndCounters map[GroupKey]*QueryInfoAndCounters
+	sync.RWMutex
 }
 
 // Reset clears the collection of statistics
 func (s *Stats) Reset() {
+	s.Lock()
+	defer s.Unlock()
+
 	s.queryInfoAndCounters = make(map[GroupKey]*QueryInfoAndCounters)
 }
 
@@ -40,7 +64,7 @@ func (s *Stats) Reset() {
 func (s *Stats) Add(doc proto.SystemProfile) error {
 	fp, err := s.fingerprinter.Fingerprint(doc.Query)
 	if err != nil {
-		return fmt.Errorf("cannot get fingerprint: %s", err)
+		return &StatsFingerprintError{err}
 	}
 	var qiac *QueryInfoAndCounters
 	var ok bool
@@ -50,8 +74,11 @@ func (s *Stats) Add(doc proto.SystemProfile) error {
 		Fingerprint: fp,
 		Namespace:   doc.Ns,
 	}
-	if qiac, ok = s.queryInfoAndCounters[key]; !ok {
-		realQuery, _ := util.GetQueryField(doc.Query)
+	if qiac, ok = s.getQueryInfoAndCounters(key); !ok {
+		realQuery, err := util.GetQueryField(doc.Query)
+		if err != nil {
+			return &StatsGetQueryFieldError{err}
+		}
 		qiac = &QueryInfoAndCounters{
 			ID:          fmt.Sprintf("%x", md5.Sum([]byte(fmt.Sprintf("%s", key)))),
 			Operation:   doc.Op,
@@ -60,7 +87,7 @@ func (s *Stats) Add(doc proto.SystemProfile) error {
 			TableScan:   false,
 			Query:       realQuery,
 		}
-		s.queryInfoAndCounters[key] = qiac
+		s.setQueryInfoAndCounters(key, qiac)
 	}
 	qiac.Count++
 	qiac.NScanned = append(qiac.NScanned, float64(doc.DocsExamined))
@@ -80,7 +107,29 @@ func (s *Stats) Add(doc proto.SystemProfile) error {
 
 // Queries returns all collected statistics
 func (s *Stats) Queries() Queries {
-	return mapToArray(s.queryInfoAndCounters)
+	s.RLock()
+	defer s.RUnlock()
+
+	queries := []QueryInfoAndCounters{}
+	for _, v := range s.queryInfoAndCounters {
+		queries = append(queries, *v)
+	}
+	return queries
+}
+
+func (s *Stats) getQueryInfoAndCounters(key GroupKey) (*QueryInfoAndCounters, bool) {
+	s.RLock()
+	defer s.RUnlock()
+
+	v, ok := s.queryInfoAndCounters[key]
+	return v, ok
+}
+
+func (s *Stats) setQueryInfoAndCounters(key GroupKey, value *QueryInfoAndCounters) {
+	s.Lock()
+	defer s.Unlock()
+
+	s.queryInfoAndCounters[key] = value
 }
 
 // Queries is a slice of MongoDB statistics
@@ -109,6 +158,25 @@ func (q Queries) CalcTotalQueriesStats(uptime int64) QueryStats {
 	return totalStats
 }
 
+type QueryInfoAndCounters struct {
+	ID          string
+	Namespace   string
+	Operation   string
+	Query       map[string]interface{}
+	Fingerprint string
+	FirstSeen   time.Time
+	LastSeen    time.Time
+	TableScan   bool
+
+	Count          int
+	BlockedTime    Times
+	LockTime       Times
+	NReturned      []float64
+	NScanned       []float64
+	QueryTime      []float64 // in milliseconds
+	ResponseLength []float64
+}
+
 // times is an array of time.Time that implements the Sorter interface
 type Times []time.Time
 
@@ -149,25 +217,6 @@ type QueryStats struct {
 	Scanned        Statistics
 }
 
-type QueryInfoAndCounters struct {
-	ID          string
-	Namespace   string
-	Operation   string
-	Query       map[string]interface{}
-	Fingerprint string
-	FirstSeen   time.Time
-	LastSeen    time.Time
-	TableScan   bool
-
-	Count          int
-	BlockedTime    Times
-	LockTime       Times
-	NReturned      []float64
-	NScanned       []float64
-	QueryTime      []float64 // in milliseconds
-	ResponseLength []float64
-}
-
 type Statistics struct {
 	Pct    float64
 	Total  float64
@@ -258,11 +307,3 @@ func calcStats(samples []float64) Statistics {
 	s.Median, _ = stats.Median(samples)
 	return s
 }
-
-func mapToArray(stats map[GroupKey]*QueryInfoAndCounters) []QueryInfoAndCounters {
-	sa := []QueryInfoAndCounters{}
-	for _, s := range stats {
-		sa = append(sa, *s)
-	}
-	return sa
-}

src/go/pt-mongodb-summary/main_test.go

@@ -5,7 +5,6 @@ import (
 	"io/ioutil"
 	"os"
 	"reflect"
-	"strings"
 	"testing"
 	"time"
 
@@ -393,29 +392,25 @@ func TestParseArgs(t *testing.T) {
 		{
 			args: []string{TOOLNAME}, // arg[0] is the command itself
 			want: &options{
-				Host:            DEFAULT_HOST,
-				LogLevel:        DEFAULT_LOGLEVEL,
-				OrderBy:         strings.Split(DEFAULT_ORDERBY, ","),
-				SkipCollections: strings.Split(DEFAULT_SKIPCOLLECTIONS, ","),
-				AuthDB:          DEFAULT_AUTHDB,
+				Host:     DEFAULT_HOST,
+				LogLevel: DEFAULT_LOGLEVEL,
+				AuthDB:   DEFAULT_AUTHDB,
 			},
 		},
 		{
 			args: []string{TOOLNAME, "zapp.brannigan.net:27018/samples", "--help"},
 			want: &options{
-				Host:            "zapp.brannigan.net:27018/samples",
-				LogLevel:        DEFAULT_LOGLEVEL,
-				OrderBy:         strings.Split(DEFAULT_ORDERBY, ","),
-				SkipCollections: strings.Split(DEFAULT_SKIPCOLLECTIONS, ","),
-				AuthDB:          DEFAULT_AUTHDB,
-				Help:            true,
+				Host:     "zapp.brannigan.net:27018/samples",
+				LogLevel: DEFAULT_LOGLEVEL,
+				AuthDB:   DEFAULT_AUTHDB,
+				Help:     true,
 			},
 		},
 	}
 	for i, test := range tests {
 		getopt.Reset()
 		os.Args = test.args
-		got, err := getOptions()
+		got, err := parseFlags()
 		if err != nil {
 			t.Errorf("error parsing command line arguments: %s", err.Error())
 		}