expvar: sort maps, fix race

It's pretty distracting to use expvar with the output of both
the top-level map and map values jumping around randomly.

Also fixes a potential race where multiple clients trying to
increment a map int or float key at the same time could lose
updates.

R=golang-codereviews, couchmoney
CC=golang-codereviews
https://golang.org/cl/54320043

diff --git a/src/pkg/expvar/expvar.go b/src/pkg/expvar/expvar.go
index b06599505fc1dccc5780c94fdcfeac67abeb445e..c590782a8d266f0fad192cccfa009d51f248a919 100644 (file)
--- a/src/pkg/expvar/expvar.go
+++ b/src/pkg/expvar/expvar.go
@@ -29,6 +29,7 @@ import (
        "net/http"
        "os"
        "runtime"
+       "sort"
        "strconv"
        "sync"
 )
@@ -40,8 +41,8 @@ type Var interface {
 
 // Int is a 64-bit integer variable that satisfies the Var interface.
 type Int struct {
-       i  int64
        mu sync.RWMutex
+       i  int64
 }
 
 func (v *Int) String() string {
@@ -64,8 +65,8 @@ func (v *Int) Set(value int64) {
 
 // Float is a 64-bit float variable that satisfies the Var interface.
 type Float struct {
-       f  float64
        mu sync.RWMutex
+       f  float64
 }
 
 func (v *Float) String() string {
@@ -90,8 +91,9 @@ func (v *Float) Set(value float64) {
 
 // Map is a string-to-Var map variable that satisfies the Var interface.
 type Map struct {
-       m  map[string]Var
-       mu sync.RWMutex
+       mu   sync.RWMutex
+       m    map[string]Var
+       keys []string // sorted
 }
 
 // KeyValue represents a single entry in a Map.
@@ -106,13 +108,13 @@ func (v *Map) String() string {
        var b bytes.Buffer
        fmt.Fprintf(&b, "{")
        first := true
-       for key, val := range v.m {
+       v.Do(func(kv KeyValue) {
                if !first {
                        fmt.Fprintf(&b, ", ")
                }
-               fmt.Fprintf(&b, "\"%s\": %v", key, val)
+               fmt.Fprintf(&b, "\"%s\": %v", kv.Key, kv.Value)
                first = false
-       }
+       })
        fmt.Fprintf(&b, "}")
        return b.String()
 }
@@ -122,6 +124,20 @@ func (v *Map) Init() *Map {
        return v
 }
 
+// updateKeys updates the sorted list of keys in v.keys.
+// must be called with v.mu held.
+func (v *Map) updateKeys() {
+       if len(v.m) == len(v.keys) {
+               // No new key.
+               return
+       }
+       v.keys = v.keys[:0]
+       for k := range v.m {
+               v.keys = append(v.keys, k)
+       }
+       sort.Strings(v.keys)
+}
+
 func (v *Map) Get(key string) Var {
        v.mu.RLock()
        defer v.mu.RUnlock()
@@ -132,6 +148,7 @@ func (v *Map) Set(key string, av Var) {
        v.mu.Lock()
        defer v.mu.Unlock()
        v.m[key] = av
+       v.updateKeys()
 }
 
 func (v *Map) Add(key string, delta int64) {
@@ -141,9 +158,11 @@ func (v *Map) Add(key string, delta int64) {
        if !ok {
                // check again under the write lock
                v.mu.Lock()
-               if _, ok = v.m[key]; !ok {
+               av, ok = v.m[key]
+               if !ok {
                        av = new(Int)
                        v.m[key] = av
+                       v.updateKeys()
                }
                v.mu.Unlock()
        }
@@ -162,9 +181,11 @@ func (v *Map) AddFloat(key string, delta float64) {
        if !ok {
                // check again under the write lock
                v.mu.Lock()
-               if _, ok = v.m[key]; !ok {
+               av, ok = v.m[key]
+               if !ok {
                        av = new(Float)
                        v.m[key] = av
+                       v.updateKeys()
                }
                v.mu.Unlock()
        }
@@ -181,15 +202,15 @@ func (v *Map) AddFloat(key string, delta float64) {
 func (v *Map) Do(f func(KeyValue)) {
        v.mu.RLock()
        defer v.mu.RUnlock()
-       for k, v := range v.m {
-               f(KeyValue{k, v})
+       for _, k := range v.keys {
+               f(KeyValue{k, v.m[k]})
        }
 }
 
 // String is a string variable, and satisfies the Var interface.
 type String struct {
-       s  string
        mu sync.RWMutex
+       s  string
 }
 
 func (v *String) String() string {
@@ -215,8 +236,9 @@ func (f Func) String() string {
 
 // All published variables.
 var (
-       mutex sync.RWMutex
-       vars  map[string]Var = make(map[string]Var)
+       mutex   sync.RWMutex
+       vars    = make(map[string]Var)
+       varKeys []string // sorted
 )
 
 // Publish declares a named exported variable. This should be called from a
@@ -229,6 +251,8 @@ func Publish(name string, v Var) {
                log.Panicln("Reuse of exported var name:", name)
        }
        vars[name] = v
+       varKeys = append(varKeys, name)
+       sort.Strings(varKeys)
 }
 
 // Get retrieves a named exported variable.
@@ -270,8 +294,8 @@ func NewString(name string) *String {
 func Do(f func(KeyValue)) {
        mutex.RLock()
        defer mutex.RUnlock()
-       for k, v := range vars {
-               f(KeyValue{k, v})
+       for _, k := range varKeys {
+               f(KeyValue{k, vars[k]})
        }
 }
 

diff --git a/src/pkg/expvar/expvar_test.go b/src/pkg/expvar/expvar_test.go
index 572c62beed56cb3eb434ff8ebc124e3512c7822b..d2ea484935e9c5fe08cb52e4c27c01a26b7cab79 100644 (file)
--- a/src/pkg/expvar/expvar_test.go
+++ b/src/pkg/expvar/expvar_test.go
@@ -5,7 +5,10 @@
 package expvar
 
 import (
+       "bytes"
        "encoding/json"
+       "net/http/httptest"
+       "strconv"
        "testing"
 )
 
@@ -15,6 +18,7 @@ func RemoveAll() {
        mutex.Lock()
        defer mutex.Unlock()
        vars = make(map[string]Var)
+       varKeys = nil
 }
 
 func TestInt(t *testing.T) {
@@ -139,3 +143,25 @@ func TestFunc(t *testing.T) {
                t.Errorf(`f.String() = %q, want %q`, s, exp)
        }
 }
+
+func TestHandler(t *testing.T) {
+       RemoveAll()
+       m := NewMap("map1")
+       m.Add("a", 1)
+       m.Add("z", 2)
+       m2 := NewMap("map2")
+       for i := 0; i < 9; i++ {
+               m2.Add(strconv.Itoa(i), int64(i))
+       }
+       rr := httptest.NewRecorder()
+       rr.Body = new(bytes.Buffer)
+       expvarHandler(rr, nil)
+       want := `{
+"map1": {"a": 1, "z": 2},
+"map2": {"0": 0, "1": 1, "2": 2, "3": 3, "4": 4, "5": 5, "6": 6, "7": 7, "8": 8}
+}
+`
+       if got := rr.Body.String(); got != want {
+               t.Errorf("HTTP handler wrote:\n%s\nWant:\n%s", got, want)
+       }
+}