├── .gitignore ├── testdata ├── fast_exit │ ├── .gitignore │ └── fast_exit.go └── empty_flush │ ├── .gitignore │ └── empty_flush.go ├── go.mod ├── package.go ├── CODE_OF_CONDUCT.md ├── go.sum ├── mock ├── sink_interface_test.go ├── example_test.go ├── sink_test.go └── sink.go ├── .github ├── dependabot.yml └── workflows │ └── actions.yml ├── logging_sink_test.go ├── .golangci.yaml ├── sink.go ├── null_sink.go ├── LICENSE ├── CONTRIBUTING.md ├── runtime_test.go ├── mock_sink.go ├── stat_handler_wrapper_1.7.go ├── stat_handler_test.go ├── stat_handler_wrapper_1.7_test.go ├── stat_handler.go ├── README.md ├── stat_handler_wrapper_test.go ├── runtime.go ├── stat_handler_wrapper.go ├── logging_sink.go ├── settings.go ├── settings_test.go ├── net_util_test.go ├── net_sink.go ├── stats_test.go ├── internal └── tags │ ├── tags.go │ └── tags_test.go ├── stats.go └── net_sink_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | vendor 3 | cover.* 4 | -------------------------------------------------------------------------------- /testdata/fast_exit/.gitignore: -------------------------------------------------------------------------------- 1 | fast_exit 2 | *.exe 3 | -------------------------------------------------------------------------------- /testdata/empty_flush/.gitignore: -------------------------------------------------------------------------------- 1 | empty_flush 2 | *.exe 3 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/lyft/gostats 2 | 3 | go 1.18 4 | 5 | require github.com/kelseyhightower/envconfig v1.4.0 6 | -------------------------------------------------------------------------------- /package.go: -------------------------------------------------------------------------------- 1 | // Package stats is a statistics library created by Engineers at Lyft 2 | // with support for Counters, Gauges, and Timers. 3 | package stats 4 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | This project is governed by [Lyft's code of conduct](https://github.com/lyft/code-of-conduct). All contributors and participants agree to abide by its terms. -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/kelseyhightower/envconfig v1.4.0 h1:Im6hONhd3pLkfDFsbRgu68RDNkGF1r3dvMUtDTo2cv8= 2 | github.com/kelseyhightower/envconfig v1.4.0/go.mod h1:cccZRl6mQpaq41TPp5QxidR+Sa3axMbJDNb//FQX6Gg= 3 | -------------------------------------------------------------------------------- /mock/sink_interface_test.go: -------------------------------------------------------------------------------- 1 | package mock_test 2 | 3 | import ( 4 | stats "github.com/lyft/gostats" 5 | "github.com/lyft/gostats/mock" 6 | ) 7 | 8 | var ( 9 | _ stats.Sink = (*mock.Sink)(nil) 10 | _ stats.FlushableSink = (*mock.Sink)(nil) 11 | ) 12 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | - package-ecosystem: github-actions 4 | directory: "/" 5 | schedule: 6 | interval: weekly 7 | open-pull-requests-limit: 99 8 | - package-ecosystem: gomod 9 | directory: "/" 10 | schedule: 11 | interval: weekly 12 | open-pull-requests-limit: 99 13 | -------------------------------------------------------------------------------- /testdata/fast_exit/fast_exit.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import stats "github.com/lyft/gostats" 4 | 5 | // Test that the stats of a fast exiting program (such as one that immediately 6 | // errors on startup) can send stats. 7 | func main() { 8 | store := stats.NewDefaultStore() 9 | store.Scope("test.fast.exit").NewCounter("counter").Inc() 10 | store.Flush() 11 | } 12 | -------------------------------------------------------------------------------- /logging_sink_test.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "os" 5 | ) 6 | 7 | func Example_flushCounter() { 8 | l := &loggingSink{writer: os.Stdout, now: foreverNow} 9 | l.FlushCounter("counterName", 420) 10 | // Output: 11 | // {"level":"debug","ts":1640995200.000000,"logger":"gostats.loggingsink","msg":"flushing counter","json":{"name":"counterName","type":"counter","value":"420.000000"}} 12 | } 13 | -------------------------------------------------------------------------------- /.golangci.yaml: -------------------------------------------------------------------------------- 1 | linters: 2 | disable-all: true 3 | enable: 4 | - gofumpt 5 | - goimports 6 | - gosimple 7 | - govet 8 | - ineffassign 9 | - misspell 10 | - nakedret 11 | - revive 12 | - staticcheck 13 | - typecheck 14 | - unconvert 15 | - unparam 16 | - unused 17 | issues: 18 | max-per-linter: 0 19 | max-same-issues: 0 20 | run: 21 | deadline: 5m 22 | -------------------------------------------------------------------------------- /sink.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | // A Sink is used by a Store to flush its data. 4 | // These functions may buffer the given data. 5 | type Sink interface { 6 | FlushCounter(name string, value uint64) 7 | FlushGauge(name string, value uint64) 8 | FlushTimer(name string, value float64) 9 | } 10 | 11 | // FlushableSink is an extension of Sink that provides a Flush() function that 12 | // will flush any buffered stats to the underlying store. 13 | type FlushableSink interface { 14 | Sink 15 | Flush() 16 | } 17 | -------------------------------------------------------------------------------- /null_sink.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | type nullSink struct{} 4 | 5 | // NewNullSink returns a Sink that does not have a backing store attached to it. 6 | func NewNullSink() FlushableSink { 7 | return nullSink{} 8 | } 9 | 10 | func (s nullSink) FlushCounter(name string, value uint64) {} //nolint:revive 11 | 12 | func (s nullSink) FlushGauge(name string, value uint64) {} //nolint:revive 13 | 14 | func (s nullSink) FlushTimer(name string, value float64) {} //nolint:revive 15 | 16 | func (s nullSink) Flush() {} 17 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | gostats 2 | 3 | Copyright 2017 Lyft Inc. 4 | 5 | Licensed under the Apache License, Version 2.0 (the "License"); 6 | you may not use this file except in compliance with the License. 7 | You may obtain a copy of the License at 8 | 9 | http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | Unless required by applicable law or agreed to in writing, software 12 | distributed under the License is distributed on an "AS IS" BASIS, 13 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | See the License for the specific language governing permissions and 15 | limitations under the License. -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | We welcome contributions from the community. Here are some guidelines. 2 | 3 | # Coding style 4 | 5 | * Gostats uses golang's `fmt` too. 6 | 7 | # Submitting a PR 8 | 9 | * Fork the repo and create your PR. 10 | * Tests will automatically run for you. 11 | * When all of the tests are passing, tag @lyft/core-libraries and @lyft/observability and we will review it and 12 | merge once our CLA has been signed (see below). 13 | * Party time. 14 | 15 | # CLA 16 | 17 | * We require a CLA for code contributions, so before we can accept a pull request we need 18 | to have a signed CLA. Please visit our [CLA service](https://oss.lyft.com/cla) and follow 19 | the instructions to sign the CLA. 20 | -------------------------------------------------------------------------------- /.github/workflows/actions.yml: -------------------------------------------------------------------------------- 1 | name: Build 2 | on: 3 | push: 4 | branches: [master] 5 | pull_request: 6 | jobs: 7 | run-tests: 8 | strategy: 9 | matrix: 10 | go-version: 11 | - 1.22.x 12 | 13 | name: run-tests 14 | runs-on: ubuntu-latest 15 | steps: 16 | - uses: actions/setup-go@v3.3.1 17 | id: go 18 | with: 19 | stable: false 20 | go-version: ${{ matrix.go-version }} 21 | 22 | - uses: actions/checkout@v3.1.0 23 | 24 | - name: run tests 25 | run: go test -vet all -race ./... 26 | 27 | lint: 28 | name: lint 29 | runs-on: ubuntu-latest 30 | steps: 31 | - uses: actions/checkout@v3.1.0 32 | 33 | - uses: golangci/golangci-lint-action@v3.3.1 34 | -------------------------------------------------------------------------------- /runtime_test.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "regexp" 5 | "testing" 6 | ) 7 | 8 | func TestRuntime(t *testing.T) { 9 | sink := &testStatSink{} 10 | store := NewStore(sink, true) 11 | 12 | scope := store.Scope("runtime") 13 | g := NewRuntimeStats(scope) 14 | g.GenerateStats() 15 | store.Flush() 16 | 17 | pattern := `runtime.lastGC:\d+|g\n` + 18 | `runtime.pauseTotalNs:\d+|g\n` + 19 | `runtime.numGC:\d+|g\n` + 20 | `runtime.alloc:\d+|g\n` + 21 | `runtime.totalAlloc:\d+|g\n` + 22 | `runtime.frees:\d+|g\n` + 23 | `runtime.nextGC:\d+|g\n` 24 | 25 | ok, err := regexp.Match(pattern, []byte(sink.record)) 26 | if err != nil { 27 | t.Fatal("Error: ", err) 28 | } 29 | if !ok { 30 | t.Errorf("Expected: '%s' Got: '%s'", pattern, sink.record) 31 | } 32 | } 33 | -------------------------------------------------------------------------------- /testdata/empty_flush/empty_flush.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "os" 7 | "time" 8 | 9 | stats "github.com/lyft/gostats" 10 | ) 11 | 12 | const FlushTimeout = time.Second * 3 13 | 14 | func flushStore(store stats.Store) error { 15 | done := make(chan struct{}) 16 | go func() { store.Flush(); close(done) }() 17 | select { 18 | case <-done: 19 | return nil 20 | case <-time.After(FlushTimeout): 21 | return errors.New("stats flush timed out") 22 | } 23 | } 24 | 25 | func realMain() (err error) { 26 | store := stats.NewDefaultStore() 27 | 28 | scope := store.ScopeWithTags("test.service.name", map[string]string{ 29 | "wrap": "1", 30 | }) 31 | defer func() { 32 | err = flushStore(store) 33 | }() 34 | _ = scope.NewCounter("panics") 35 | return 36 | } 37 | 38 | func main() { 39 | if err := realMain(); err != nil { 40 | fmt.Fprintln(os.Stderr, "Error:", err) 41 | os.Exit(1) 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /mock_sink.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import "sync" 4 | 5 | // MockSink describes an in-memory Sink used for testing. 6 | // 7 | // DEPRECATED: use "github.com/lyft/gostats/mock" instead. 8 | type MockSink struct { 9 | Counters map[string]uint64 10 | Timers map[string]uint64 11 | Gauges map[string]uint64 12 | 13 | cLock sync.Mutex 14 | tLock sync.Mutex 15 | gLock sync.Mutex 16 | } 17 | 18 | // NewMockSink returns a MockSink that flushes stats to in-memory maps. An 19 | // instance of MockSink is not safe for concurrent use. 20 | // 21 | // DEPRECATED: use "github.com/lyft/gostats/mock" instead. 22 | func NewMockSink() (m *MockSink) { 23 | m = &MockSink{ 24 | Counters: make(map[string]uint64), 25 | Timers: make(map[string]uint64), 26 | Gauges: make(map[string]uint64), 27 | } 28 | 29 | return 30 | } 31 | 32 | // FlushCounter satisfies the Sink interface. 33 | func (m *MockSink) FlushCounter(name string, value uint64) { 34 | m.cLock.Lock() 35 | defer m.cLock.Unlock() 36 | m.Counters[name] += value 37 | } 38 | 39 | // FlushGauge satisfies the Sink interface. 40 | func (m *MockSink) FlushGauge(name string, value uint64) { 41 | m.gLock.Lock() 42 | defer m.gLock.Unlock() 43 | m.Gauges[name] = value 44 | } 45 | 46 | // FlushTimer satisfies the Sink interface. 47 | func (m *MockSink) FlushTimer(name string, value float64) { //nolint:revive 48 | m.tLock.Lock() 49 | defer m.tLock.Unlock() 50 | m.Timers[name]++ 51 | } 52 | -------------------------------------------------------------------------------- /stat_handler_wrapper_1.7.go: -------------------------------------------------------------------------------- 1 | //go:build !go1.8 2 | // +build !go1.8 3 | 4 | package stats 5 | 6 | import "net/http" 7 | 8 | func (h *httpHandler) wrapResponse(w http.ResponseWriter) http.ResponseWriter { 9 | rw := &responseWriter{ 10 | ResponseWriter: w, 11 | handler: h, 12 | } 13 | 14 | flusher, canFlush := w.(http.Flusher) 15 | hijacker, canHijack := w.(http.Hijacker) 16 | closeNotifier, canNotify := w.(http.CloseNotifier) 17 | 18 | if canFlush && canHijack && canNotify { 19 | return struct { 20 | http.ResponseWriter 21 | http.Flusher 22 | http.Hijacker 23 | http.CloseNotifier 24 | }{rw, flusher, hijacker, closeNotifier} 25 | } else if canFlush && canHijack { 26 | return struct { 27 | http.ResponseWriter 28 | http.Flusher 29 | http.Hijacker 30 | }{rw, flusher, hijacker} 31 | } else if canFlush && canNotify { 32 | return struct { 33 | http.ResponseWriter 34 | http.Flusher 35 | http.CloseNotifier 36 | }{rw, flusher, closeNotifier} 37 | } else if canHijack && canNotify { 38 | return struct { 39 | http.ResponseWriter 40 | http.Hijacker 41 | http.CloseNotifier 42 | }{rw, hijacker, closeNotifier} 43 | } else if canFlush { 44 | return struct { 45 | http.ResponseWriter 46 | http.Flusher 47 | }{rw, flusher} 48 | } else if canHijack { 49 | return struct { 50 | http.ResponseWriter 51 | http.Hijacker 52 | }{rw, hijacker} 53 | } else if canNotify { 54 | return struct { 55 | http.ResponseWriter 56 | http.CloseNotifier 57 | }{rw, closeNotifier} 58 | } 59 | 60 | return rw 61 | } 62 | -------------------------------------------------------------------------------- /stat_handler_test.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "io" 5 | "net/http" 6 | "net/http/httptest" 7 | "strconv" 8 | "strings" 9 | "sync" 10 | "testing" 11 | 12 | "github.com/lyft/gostats/mock" 13 | ) 14 | 15 | func TestHttpHandler_ServeHTTP(t *testing.T) { 16 | t.Parallel() 17 | 18 | sink := mock.NewSink() 19 | store := NewStore(sink, false) 20 | 21 | h := NewStatHandler( 22 | store, 23 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 24 | if code, err := strconv.Atoi(r.Header.Get("code")); err == nil { 25 | w.WriteHeader(code) 26 | } 27 | 28 | io.Copy(w, r.Body) 29 | r.Body.Close() 30 | })).(*httpHandler) 31 | 32 | wg := sync.WaitGroup{} 33 | wg.Add(2) 34 | 35 | go func() { 36 | r, _ := http.NewRequest(http.MethodGet, "/", strings.NewReader("foo")) 37 | w := httptest.NewRecorder() 38 | h.ServeHTTP(w, r) 39 | store.Flush() 40 | 41 | if w.Body.String() != "foo" { 42 | t.Errorf("wanted %q body, got %q", "foo", w.Body.String()) 43 | } 44 | 45 | if w.Code != http.StatusOK { 46 | t.Errorf("wanted 200, got %d", w.Code) 47 | } 48 | 49 | wg.Done() 50 | }() 51 | 52 | go func() { 53 | r := httptest.NewRequest(http.MethodGet, "/", strings.NewReader("bar")) 54 | r.Header.Set("code", strconv.Itoa(http.StatusNotFound)) 55 | w := httptest.NewRecorder() 56 | h.ServeHTTP(w, r) 57 | store.Flush() 58 | 59 | if w.Body.String() != "bar" { 60 | t.Errorf("wanted %q body, got %q", "bar", w.Body.String()) 61 | } 62 | 63 | if w.Code != http.StatusNotFound { 64 | t.Errorf("wanted 404, got %d", w.Code) 65 | } 66 | 67 | wg.Done() 68 | }() 69 | 70 | wg.Wait() 71 | 72 | sink.AssertTimerCallCount(t, requestTimer, 2) 73 | sink.AssertCounterEquals(t, "200", 1) 74 | sink.AssertCounterEquals(t, "404", 1) 75 | } 76 | -------------------------------------------------------------------------------- /mock/example_test.go: -------------------------------------------------------------------------------- 1 | package mock_test 2 | 3 | import ( 4 | "fmt" 5 | "reflect" 6 | "testing" 7 | 8 | stats "github.com/lyft/gostats" 9 | "github.com/lyft/gostats/mock" 10 | ) 11 | 12 | func ExampleSerializeTags() { 13 | sink := mock.NewSink() 14 | store := stats.NewStore(sink, false) 15 | 16 | tags := map[string]string{ 17 | "key_1": "val_1", 18 | "key_2": "val_2", 19 | } 20 | counter := store.NewCounterWithTags("counter", tags) 21 | counter.Add(2) 22 | 23 | store.Flush() 24 | 25 | n := sink.Counter(mock.SerializeTags("counter", tags)) 26 | fmt.Println("counter:", n, n == 2) 27 | 28 | // Output: 29 | // counter: 2 true 30 | } 31 | 32 | func ExampleParseTags() { 33 | sink := mock.NewSink() 34 | store := stats.NewStore(sink, false) 35 | 36 | tags := map[string]string{ 37 | "key_1": "val_1", 38 | "key_2": "val_2", 39 | } 40 | for i := 0; i < 4; i++ { 41 | store.NewCounterWithTags(fmt.Sprintf("c_%d", i), tags).Inc() 42 | } 43 | store.Flush() 44 | 45 | // Check that the counters all have the same tags 46 | for stat := range sink.Counters() { 47 | name, m := mock.ParseTags(stat) 48 | if !reflect.DeepEqual(m, tags) { 49 | panic(fmt.Sprintf("Tags: got: %q want: %q", m, tags)) 50 | } 51 | fmt.Printf("%s: okay\n", name) 52 | } 53 | 54 | // Unordered output: 55 | // c_0: okay 56 | // c_1: okay 57 | // c_2: okay 58 | // c_3: okay 59 | } 60 | 61 | func ExampleFatal() { 62 | sink := mock.NewSink() 63 | store := stats.NewStore(sink, false) 64 | 65 | store.NewCounter("c").Set(2) 66 | store.Flush() 67 | 68 | // In real test code you would use the *testing.T or *testing.B 69 | // passed to the test function. 70 | t := &testing.T{} 71 | 72 | // This will cause the test to immediately fail with .Fatal() 73 | // if the assertion is false. 74 | sink.AssertCounterEquals(mock.Fatal(t), "c", 2) 75 | } 76 | -------------------------------------------------------------------------------- /stat_handler_wrapper_1.7_test.go: -------------------------------------------------------------------------------- 1 | //go:build !go1.8 2 | // +build !go1.8 3 | 4 | package stats 5 | 6 | import ( 7 | "fmt" 8 | "net/http" 9 | "testing" 10 | ) 11 | 12 | func TestHTTPHandler_WrapResponse(t *testing.T) { 13 | t.Parallel() 14 | 15 | tests := []http.ResponseWriter{ 16 | struct { 17 | http.ResponseWriter 18 | http.Flusher 19 | http.Hijacker 20 | http.CloseNotifier 21 | }{}, 22 | struct { 23 | http.ResponseWriter 24 | http.Flusher 25 | http.Hijacker 26 | }{}, 27 | struct { 28 | http.ResponseWriter 29 | http.Flusher 30 | http.CloseNotifier 31 | }{}, 32 | struct { 33 | http.ResponseWriter 34 | http.Flusher 35 | }{}, 36 | struct { 37 | http.ResponseWriter 38 | http.Hijacker 39 | http.CloseNotifier 40 | }{}, 41 | struct { 42 | http.ResponseWriter 43 | http.Hijacker 44 | }{}, 45 | struct { 46 | http.ResponseWriter 47 | http.CloseNotifier 48 | }{}, 49 | struct { 50 | http.ResponseWriter 51 | }{}, 52 | } 53 | 54 | h := NewStatHandler( 55 | NewStore(NewNullSink(), false), 56 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})).(*httpHandler) 57 | 58 | for i, test := range tests { 59 | tc := test 60 | t.Run(fmt.Sprint("test:", i), func(t *testing.T) { 61 | t.Parallel() 62 | 63 | _, canFlush := tc.(http.Flusher) 64 | _, canHijack := tc.(http.Hijacker) 65 | _, canNotify := tc.(http.CloseNotifier) 66 | 67 | rw := h.wrapResponse(tc) 68 | 69 | if _, ok := rw.(http.Flusher); ok != canFlush { 70 | t.Errorf("Flusher: wanted %t", canFlush) 71 | } 72 | if _, ok := rw.(http.Hijacker); ok != canHijack { 73 | t.Errorf("Hijacker: wanted %t", canHijack) 74 | } 75 | if _, ok := rw.(http.CloseNotifier); ok != canNotify { 76 | t.Errorf("CloseNotifier: wanted %t", canNotify) 77 | } 78 | }) 79 | } 80 | } 81 | -------------------------------------------------------------------------------- /stat_handler.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "net/http" 5 | "strconv" 6 | "sync" 7 | ) 8 | 9 | const requestTimer = "rq_time_us" 10 | 11 | type httpHandler struct { 12 | scope Scope 13 | delegate http.Handler 14 | 15 | timer Timer 16 | 17 | codes map[int]Counter 18 | codesMtx sync.RWMutex 19 | } 20 | 21 | // NewStatHandler returns an http handler for stats. 22 | func NewStatHandler(scope Scope, handler http.Handler) http.Handler { 23 | return &httpHandler{ 24 | scope: scope, 25 | delegate: handler, 26 | timer: scope.NewTimer(requestTimer), 27 | codes: map[int]Counter{}, 28 | } 29 | } 30 | 31 | func (h *httpHandler) counter(code int) Counter { 32 | h.codesMtx.RLock() 33 | c := h.codes[code] 34 | h.codesMtx.RUnlock() 35 | 36 | if c != nil { 37 | return c 38 | } 39 | 40 | h.codesMtx.Lock() 41 | if c = h.codes[code]; c == nil { 42 | c = h.scope.NewCounter(strconv.Itoa(code)) 43 | h.codes[code] = c 44 | } 45 | h.codesMtx.Unlock() 46 | 47 | return c 48 | } 49 | 50 | func (h *httpHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { 51 | span := h.timer.AllocateSpan() 52 | h.delegate.ServeHTTP(h.wrapResponse(w), r) 53 | span.Complete() 54 | } 55 | 56 | type responseWriter struct { 57 | http.ResponseWriter 58 | 59 | headerWritten bool 60 | handler *httpHandler 61 | } 62 | 63 | func (rw *responseWriter) Write(b []byte) (int, error) { 64 | if !rw.headerWritten { 65 | rw.WriteHeader(http.StatusOK) 66 | } 67 | return rw.ResponseWriter.Write(b) 68 | } 69 | 70 | func (rw *responseWriter) WriteHeader(code int) { 71 | if rw.headerWritten { 72 | return 73 | } 74 | 75 | rw.headerWritten = true 76 | rw.handler.counter(code).Inc() 77 | rw.ResponseWriter.WriteHeader(code) 78 | } 79 | 80 | var ( 81 | _ http.Handler = (*httpHandler)(nil) 82 | _ http.ResponseWriter = (*responseWriter)(nil) 83 | ) 84 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Gostats [![GoDoc](https://godoc.org/github.com/lyft/gostats?status.svg)](https://godoc.org/github.com/lyft/gostats) [![Build Status](https://github.com/lyft/gostats/actions/workflows/actions.yml/badge.svg?branch=master)](https://github.com/lyft/gostats/actions/workflows/actions.yml) 2 | 3 | `gostats` is a Go metrics library with support for Counters, Gauges, and Timers. 4 | 5 | ## Installation 6 | 7 | ```sh 8 | go get github.com/lyft/gostats 9 | ``` 10 | 11 | ## Building & Testing 12 | 13 | ```sh 14 | go test ./... 15 | ``` 16 | 17 | ## Usage 18 | 19 | In order to start using `gostats`, import it into your project with: 20 | 21 | ```go 22 | import "github.com/lyft/gostats" 23 | ``` 24 | 25 | 26 | ## Mocking 27 | 28 | A thread-safe mock sink is provided by the [gostats/mock](https://github.com/lyft/gostats/blob/mock-sink/mock/sink.go) package. The mock sink also provides methods that are useful for testing (as demonstrated below). 29 | ```go 30 | package mock_test 31 | 32 | import ( 33 | "testing" 34 | 35 | "github.com/lyft/gostats" 36 | "github.com/lyft/gostats/mock" 37 | ) 38 | 39 | type Config struct { 40 | Stats stats.Store 41 | } 42 | 43 | func TestMockExample(t *testing.T) { 44 | sink := mock.NewSink() 45 | conf := Config{ 46 | Stats: stats.NewStore(sink, false), 47 | } 48 | conf.Stats.NewCounter("name").Inc() 49 | conf.Stats.Flush() 50 | sink.AssertCounterEquals(t, "name", 1) 51 | } 52 | ``` 53 | 54 | If you do not need to assert on the contents of the sink the below example can be used to quickly create a thread-safe `stats.Scope`: 55 | ```go 56 | package config 57 | 58 | import ( 59 | "github.com/lyft/gostats" 60 | "github.com/lyft/gostats/mock" 61 | ) 62 | 63 | type Config struct { 64 | Stats stats.Store 65 | } 66 | 67 | func NewConfig() *Config { 68 | return &Config{ 69 | Stats: stats.NewDefaultStore(), 70 | } 71 | } 72 | 73 | func NewMockConfig() *Config { 74 | sink := mock.NewSink() 75 | return &Config{ 76 | Stats: stats.NewStore(sink, false), 77 | } 78 | } 79 | ``` 80 | -------------------------------------------------------------------------------- /stat_handler_wrapper_test.go: -------------------------------------------------------------------------------- 1 | //go:build go1.8 2 | // +build go1.8 3 | 4 | package stats 5 | 6 | import ( 7 | "net/http" 8 | "testing" 9 | ) 10 | 11 | func TestHTTPHandler_WrapResponse(t *testing.T) { 12 | tests := []http.ResponseWriter{ 13 | struct { 14 | http.ResponseWriter 15 | http.Flusher 16 | http.Hijacker 17 | http.Pusher 18 | http.CloseNotifier 19 | }{}, 20 | struct { 21 | http.ResponseWriter 22 | http.Flusher 23 | http.Hijacker 24 | http.Pusher 25 | }{}, 26 | struct { 27 | http.ResponseWriter 28 | http.Flusher 29 | http.Hijacker 30 | http.CloseNotifier 31 | }{}, 32 | struct { 33 | http.ResponseWriter 34 | http.Flusher 35 | http.Hijacker 36 | }{}, 37 | struct { 38 | http.ResponseWriter 39 | http.Flusher 40 | http.Pusher 41 | http.CloseNotifier 42 | }{}, 43 | struct { 44 | http.ResponseWriter 45 | http.Flusher 46 | http.Pusher 47 | }{}, 48 | struct { 49 | http.ResponseWriter 50 | http.Flusher 51 | http.CloseNotifier 52 | }{}, 53 | struct { 54 | http.ResponseWriter 55 | http.Flusher 56 | }{}, 57 | struct { 58 | http.ResponseWriter 59 | http.Hijacker 60 | http.Pusher 61 | http.CloseNotifier 62 | }{}, 63 | struct { 64 | http.ResponseWriter 65 | http.Hijacker 66 | http.Pusher 67 | }{}, 68 | struct { 69 | http.ResponseWriter 70 | http.Hijacker 71 | http.CloseNotifier 72 | }{}, 73 | struct { 74 | http.ResponseWriter 75 | http.Hijacker 76 | }{}, 77 | struct { 78 | http.ResponseWriter 79 | http.Pusher 80 | http.CloseNotifier 81 | }{}, 82 | struct { 83 | http.ResponseWriter 84 | http.Pusher 85 | }{}, 86 | struct { 87 | http.ResponseWriter 88 | http.CloseNotifier 89 | }{}, 90 | struct{ http.ResponseWriter }{}, 91 | } 92 | 93 | h := NewStatHandler( 94 | NewStore(NewNullSink(), false), 95 | http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})).(*httpHandler) 96 | 97 | for i, test := range tests { 98 | tc := test 99 | _, canFlush := tc.(http.Flusher) 100 | _, canHijack := tc.(http.Hijacker) 101 | _, canPush := tc.(http.Pusher) 102 | 103 | rw := h.wrapResponse(tc) 104 | 105 | if _, ok := rw.(http.Flusher); ok != canFlush { 106 | t.Errorf("Test(%d): Flusher: wanted %t", i, canFlush) 107 | } 108 | if _, ok := rw.(http.Hijacker); ok != canHijack { 109 | t.Errorf("Test(%d): Hijacker: wanted %t", i, canHijack) 110 | } 111 | if _, ok := rw.(http.Pusher); ok != canPush { 112 | t.Errorf("Test(%d): Pusher: wanted %t", i, canPush) 113 | } 114 | } 115 | } 116 | -------------------------------------------------------------------------------- /runtime.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "runtime" 5 | ) 6 | 7 | type runtimeStats struct { 8 | alloc Gauge // bytes allocated and not yet freed 9 | totalAlloc Counter // bytes allocated (even if freed) 10 | sys Gauge // bytes obtained from system (sum of XxxSys below) 11 | lookups Counter // number of pointer lookups 12 | mallocs Counter // number of mallocs 13 | frees Counter // number of frees 14 | 15 | // Main allocation heap statistics 16 | heapAlloc Gauge // bytes allocated and not yet freed (same as Alloc above) 17 | heapSys Gauge // bytes obtained from system 18 | heapIdle Gauge // bytes in idle spans 19 | heapInuse Gauge // bytes in non-idle span 20 | heapReleased Gauge // bytes released to the OS 21 | heapObjects Gauge // total number of allocated objects 22 | 23 | // Garbage collector statistics. 24 | nextGC Gauge // next collection will happen when HeapAlloc ≥ this amount 25 | lastGC Gauge // end time of last collection (nanoseconds since 1970) 26 | pauseTotalNs Counter 27 | numGC Counter 28 | gcCPUPercent Gauge 29 | 30 | numGoroutine Gauge 31 | } 32 | 33 | // NewRuntimeStats returns a StatGenerator with common Go runtime stats like memory allocated, 34 | // total mallocs, total frees, etc. 35 | func NewRuntimeStats(scope Scope) StatGenerator { 36 | return runtimeStats{ 37 | alloc: scope.NewGauge("alloc"), 38 | totalAlloc: scope.NewCounter("totalAlloc"), 39 | sys: scope.NewGauge("sys"), 40 | lookups: scope.NewCounter("lookups"), 41 | mallocs: scope.NewCounter("mallocs"), 42 | frees: scope.NewCounter("frees"), 43 | 44 | heapAlloc: scope.NewGauge("heapAlloc"), 45 | heapSys: scope.NewGauge("heapSys"), 46 | heapIdle: scope.NewGauge("heapIdle"), 47 | heapInuse: scope.NewGauge("heapInuse"), 48 | heapReleased: scope.NewGauge("heapReleased"), 49 | heapObjects: scope.NewGauge("heapObjects"), 50 | 51 | nextGC: scope.NewGauge("nextGC"), 52 | lastGC: scope.NewGauge("lastGC"), 53 | pauseTotalNs: scope.NewCounter("pauseTotalNs"), 54 | numGC: scope.NewCounter("numGC"), 55 | gcCPUPercent: scope.NewGauge("gcCPUPercent"), 56 | 57 | numGoroutine: scope.NewGauge("numGoroutine"), 58 | } 59 | } 60 | 61 | func (r runtimeStats) GenerateStats() { 62 | var memStats runtime.MemStats 63 | runtime.ReadMemStats(&memStats) 64 | 65 | r.alloc.Set(memStats.Alloc) 66 | r.totalAlloc.Set(memStats.TotalAlloc) 67 | r.mallocs.Set(memStats.Mallocs) 68 | r.frees.Set(memStats.Frees) 69 | 70 | r.heapAlloc.Set(memStats.HeapAlloc) 71 | r.heapSys.Set(memStats.HeapSys) 72 | r.heapIdle.Set(memStats.HeapIdle) 73 | r.heapInuse.Set(memStats.HeapInuse) 74 | r.heapReleased.Set(memStats.HeapReleased) 75 | r.heapObjects.Set(memStats.HeapObjects) 76 | 77 | r.nextGC.Set(memStats.NextGC) 78 | r.lastGC.Set(memStats.LastGC) 79 | r.pauseTotalNs.Set(memStats.PauseTotalNs) 80 | r.numGC.Set(uint64(memStats.NumGC)) 81 | r.gcCPUPercent.Set(uint64(memStats.GCCPUFraction * 100)) 82 | 83 | r.numGoroutine.Set(uint64(runtime.NumGoroutine())) 84 | } 85 | -------------------------------------------------------------------------------- /stat_handler_wrapper.go: -------------------------------------------------------------------------------- 1 | //go:build go1.8 2 | // +build go1.8 3 | 4 | package stats 5 | 6 | import "net/http" 7 | 8 | func (h *httpHandler) wrapResponse(w http.ResponseWriter) http.ResponseWriter { 9 | rw := &responseWriter{ 10 | ResponseWriter: w, 11 | handler: h, 12 | } 13 | 14 | flusher, canFlush := w.(http.Flusher) 15 | hijacker, canHijack := w.(http.Hijacker) 16 | pusher, canPush := w.(http.Pusher) 17 | 18 | //nolint:staticcheck 19 | closeNotifier, canNotify := w.(http.CloseNotifier) 20 | 21 | if canFlush && canHijack && canPush && canNotify { 22 | return struct { 23 | http.ResponseWriter 24 | http.Flusher 25 | http.Hijacker 26 | http.Pusher 27 | http.CloseNotifier 28 | }{rw, flusher, hijacker, pusher, closeNotifier} 29 | } else if canFlush && canHijack && canPush { 30 | return struct { 31 | http.ResponseWriter 32 | http.Flusher 33 | http.Hijacker 34 | http.Pusher 35 | }{rw, flusher, hijacker, pusher} 36 | } else if canFlush && canHijack && canNotify { 37 | return struct { 38 | http.ResponseWriter 39 | http.Flusher 40 | http.Hijacker 41 | http.CloseNotifier 42 | }{rw, flusher, hijacker, closeNotifier} 43 | } else if canFlush && canPush && canNotify { 44 | return struct { 45 | http.ResponseWriter 46 | http.Flusher 47 | http.Pusher 48 | http.CloseNotifier 49 | }{rw, flusher, pusher, closeNotifier} 50 | } else if canHijack && canPush && canNotify { 51 | return struct { 52 | http.ResponseWriter 53 | http.Hijacker 54 | http.Pusher 55 | http.CloseNotifier 56 | }{rw, hijacker, pusher, closeNotifier} 57 | } else if canFlush && canHijack { 58 | return struct { 59 | http.ResponseWriter 60 | http.Flusher 61 | http.Hijacker 62 | }{rw, flusher, hijacker} 63 | } else if canFlush && canPush { 64 | return struct { 65 | http.ResponseWriter 66 | http.Flusher 67 | http.Pusher 68 | }{rw, flusher, pusher} 69 | } else if canFlush && canNotify { 70 | return struct { 71 | http.ResponseWriter 72 | http.Flusher 73 | http.CloseNotifier 74 | }{rw, flusher, closeNotifier} 75 | } else if canHijack && canPush { 76 | return struct { 77 | http.ResponseWriter 78 | http.Hijacker 79 | http.Pusher 80 | }{rw, hijacker, pusher} 81 | } else if canHijack && canNotify { 82 | return struct { 83 | http.ResponseWriter 84 | http.Hijacker 85 | http.CloseNotifier 86 | }{rw, hijacker, closeNotifier} 87 | } else if canPush && canNotify { 88 | return struct { 89 | http.ResponseWriter 90 | http.Pusher 91 | http.CloseNotifier 92 | }{rw, pusher, closeNotifier} 93 | } else if canFlush { 94 | return struct { 95 | http.ResponseWriter 96 | http.Flusher 97 | }{rw, flusher} 98 | } else if canHijack { 99 | return struct { 100 | http.ResponseWriter 101 | http.Hijacker 102 | }{rw, hijacker} 103 | } else if canPush { 104 | return struct { 105 | http.ResponseWriter 106 | http.Pusher 107 | }{rw, pusher} 108 | } else if canNotify { 109 | return struct { 110 | http.ResponseWriter 111 | http.CloseNotifier 112 | }{rw, closeNotifier} 113 | } 114 | 115 | return rw 116 | } 117 | -------------------------------------------------------------------------------- /logging_sink.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "io" 7 | "os" 8 | "strconv" 9 | "time" 10 | ) 11 | 12 | type loggingSink struct { 13 | writer io.Writer 14 | now func() time.Time 15 | } 16 | 17 | type logLine struct { 18 | Level string `json:"level"` 19 | Timestamp sixDecimalPlacesFloat `json:"ts"` 20 | Logger string `json:"logger"` 21 | Message string `json:"msg"` 22 | JSON map[string]string `json:"json"` 23 | } 24 | 25 | type sixDecimalPlacesFloat float64 26 | 27 | func (f sixDecimalPlacesFloat) MarshalJSON() ([]byte, error) { 28 | var ret []byte 29 | ret = strconv.AppendFloat(ret, float64(f), 'f', 6, 64) 30 | return ret, nil 31 | } 32 | 33 | // NewLoggingSink returns a "default" logging Sink that flushes stats 34 | // to os.StdErr. This sink is not fast, or flexible, it doesn't 35 | // buffer, it exists merely to be convenient to use by default, with 36 | // no configuration. 37 | // 38 | // The format of this logger is similar to Zap, but not explicitly 39 | // importing Zap to avoid the dependency. The format is as if you used 40 | // a zap.NewProduction-generated logger, but also added a 41 | // log.With(zap.Namespace("json")). This does not include any 42 | // stacktrace for errors at the moment. 43 | // 44 | // If these defaults do not work for you, users should provide their 45 | // own logger, conforming to FlushableSink, instead. 46 | func NewLoggingSink() FlushableSink { 47 | return &loggingSink{writer: os.Stderr, now: time.Now} 48 | } 49 | 50 | // this is allocated outside of logMessage, even though its only used 51 | // there, to avoid allocing a map every time we log. 52 | var emptyMap = map[string]string{} 53 | 54 | func (s *loggingSink) logMessage(level string, msg string) { 55 | nanos := s.now().UnixNano() 56 | sec := sixDecimalPlacesFloat(float64(nanos) / float64(time.Second)) 57 | enc := json.NewEncoder(s.writer) 58 | enc.Encode(logLine{ 59 | Message: msg, 60 | Level: level, 61 | Timestamp: sec, 62 | Logger: "gostats.loggingsink", 63 | // intentional empty map used to avoid any null parsing issues 64 | // on the log collection side 65 | JSON: emptyMap, 66 | }) 67 | } 68 | 69 | func (s *loggingSink) log(name, typ string, value float64) { 70 | nanos := s.now().UnixNano() 71 | sec := sixDecimalPlacesFloat(float64(nanos) / float64(time.Second)) 72 | enc := json.NewEncoder(s.writer) 73 | kv := map[string]string{ 74 | "type": typ, 75 | "value": fmt.Sprintf("%f", value), 76 | } 77 | if name != "" { 78 | kv["name"] = name 79 | } 80 | enc.Encode(logLine{ 81 | Message: fmt.Sprintf("flushing %s", typ), 82 | Level: "debug", 83 | Timestamp: sec, 84 | Logger: "gostats.loggingsink", 85 | JSON: kv, 86 | }) 87 | } 88 | 89 | func (s *loggingSink) FlushCounter(name string, value uint64) { s.log(name, "counter", float64(value)) } 90 | 91 | func (s *loggingSink) FlushGauge(name string, value uint64) { s.log(name, "gauge", float64(value)) } 92 | 93 | func (s *loggingSink) FlushTimer(name string, value float64) { s.log(name, "timer", value) } 94 | 95 | func (s *loggingSink) Flush() { s.log("", "all stats", 0) } 96 | 97 | // Logger 98 | 99 | func (s *loggingSink) Errorf(msg string, args ...interface{}) { 100 | s.logMessage("error", fmt.Sprintf(msg, args...)) 101 | } 102 | 103 | func (s *loggingSink) Warnf(msg string, args ...interface{}) { 104 | s.logMessage("warn", fmt.Sprintf(msg, args...)) 105 | } 106 | -------------------------------------------------------------------------------- /settings.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "strconv" 7 | "time" 8 | ) 9 | 10 | const ( 11 | // DefaultUseStatsd use statsd as a stats sink, default is true. 12 | DefaultUseStatsd = true 13 | // DefaultStatsdHost is the default address where statsd is running at. 14 | DefaultStatsdHost = "localhost" 15 | // DefaultStatsdProtocol is TCP 16 | DefaultStatsdProtocol = "tcp" 17 | // DefaultStatsdPort is the default port where statsd is listening at. 18 | DefaultStatsdPort = 8125 19 | // DefaultFlushIntervalS is the default flushing interval in seconds. 20 | DefaultFlushIntervalS = 5 21 | // DefaultLoggingSinkDisabled is the default behavior of logging sink suppression, default is false. 22 | DefaultLoggingSinkDisabled = false 23 | ) 24 | 25 | // The Settings type is used to configure gostats. gostats uses environment 26 | // variables to setup its settings. 27 | type Settings struct { 28 | // Use statsd as a stats sink. 29 | UseStatsd bool `envconfig:"USE_STATSD" default:"true"` 30 | // Address where statsd is running at. 31 | StatsdHost string `envconfig:"STATSD_HOST" default:"localhost"` 32 | // Network protocol used to connect to statsd 33 | StatsdProtocol string `envconfig:"STATSD_PROTOCOL" default:"tcp"` 34 | // Port where statsd is listening at. 35 | StatsdPort int `envconfig:"STATSD_PORT" default:"8125"` 36 | // Flushing interval. 37 | FlushIntervalS int `envconfig:"GOSTATS_FLUSH_INTERVAL_SECONDS" default:"5"` 38 | // Disable the LoggingSink when USE_STATSD is false and use the NullSink instead. 39 | // This will cause all stats to be silently dropped. 40 | LoggingSinkDisabled bool `envconfig:"GOSTATS_LOGGING_SINK_DISABLED" default:"false"` 41 | } 42 | 43 | // An envError is an error that occurred parsing an environment variable 44 | type envError struct { 45 | Key string 46 | Value string 47 | Err error 48 | } 49 | 50 | func (e *envError) Error() string { 51 | return fmt.Sprintf("parsing environment variable: %q with value: %q: %s", 52 | e.Key, e.Value, e.Err) 53 | } 54 | 55 | func envOr(key, def string) string { 56 | if s := os.Getenv(key); s != "" { 57 | return s 58 | } 59 | return def 60 | } 61 | 62 | func envInt(key string, def int) (int, error) { 63 | s := os.Getenv(key) 64 | if s == "" { 65 | return def, nil 66 | } 67 | i, err := strconv.Atoi(s) 68 | if err != nil { 69 | return def, &envError{Key: key, Value: s, Err: err} 70 | } 71 | return i, nil 72 | } 73 | 74 | func envBool(key string, def bool) (bool, error) { 75 | s := os.Getenv(key) 76 | if s == "" { 77 | return def, nil 78 | } 79 | b, err := strconv.ParseBool(s) 80 | if err != nil { 81 | return def, &envError{Key: key, Value: s, Err: err} 82 | } 83 | return b, nil 84 | } 85 | 86 | // GetSettings returns the Settings gostats will run with. 87 | func GetSettings() Settings { 88 | useStatsd, err := envBool("USE_STATSD", DefaultUseStatsd) 89 | if err != nil { 90 | panic(err) 91 | } 92 | statsdPort, err := envInt("STATSD_PORT", DefaultStatsdPort) 93 | if err != nil { 94 | panic(err) 95 | } 96 | flushIntervalS, err := envInt("GOSTATS_FLUSH_INTERVAL_SECONDS", DefaultFlushIntervalS) 97 | if err != nil { 98 | panic(err) 99 | } 100 | loggingSinkDisabled, err := envBool("GOSTATS_LOGGING_SINK_DISABLED", DefaultLoggingSinkDisabled) 101 | if err != nil { 102 | panic(err) 103 | } 104 | return Settings{ 105 | UseStatsd: useStatsd, 106 | StatsdHost: envOr("STATSD_HOST", DefaultStatsdHost), 107 | StatsdProtocol: envOr("STATSD_PROTOCOL", DefaultStatsdProtocol), 108 | StatsdPort: statsdPort, 109 | FlushIntervalS: flushIntervalS, 110 | LoggingSinkDisabled: loggingSinkDisabled, 111 | } 112 | } 113 | 114 | // FlushInterval returns the flush interval duration. 115 | func (s *Settings) FlushInterval() time.Duration { 116 | return time.Duration(s.FlushIntervalS) * time.Second 117 | } 118 | -------------------------------------------------------------------------------- /settings_test.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "os" 5 | "reflect" 6 | "testing" 7 | "time" 8 | 9 | "github.com/kelseyhightower/envconfig" 10 | ) 11 | 12 | func testSetenv(t *testing.T, pairs ...string) (reset func()) { 13 | var fns []func() 14 | for i := 0; i < len(pairs); i += 2 { 15 | key := pairs[i+0] 16 | val := pairs[i+1] 17 | 18 | prev, exists := os.LookupEnv(key) 19 | if val == "" { 20 | if err := os.Unsetenv(key); err != nil { 21 | t.Fatalf("deleting env key: %s: %s", key, err) 22 | } 23 | } else { 24 | if err := os.Setenv(key, val); err != nil { 25 | t.Fatalf("setting env key: %s: %s", key, err) 26 | } 27 | } 28 | if exists { 29 | fns = append(fns, func() { os.Setenv(key, prev) }) 30 | } else { 31 | fns = append(fns, func() { os.Unsetenv(key) }) 32 | } 33 | } 34 | return func() { 35 | for _, fn := range fns { 36 | fn() 37 | } 38 | } 39 | } 40 | 41 | func TestSettingsCompat(t *testing.T) { 42 | reset := testSetenv(t, 43 | "USE_STATSD", "", 44 | "STATSD_HOST", "", 45 | "STATSD_PROTOCOL", "", 46 | "STATSD_PORT", "", 47 | "GOSTATS_FLUSH_INTERVAL_SECONDS", "", 48 | "GOSTATS_LOGGING_SINK_DISABLED", "", 49 | ) 50 | defer reset() 51 | 52 | var e Settings 53 | if err := envconfig.Process("", &e); err != nil { 54 | t.Fatal(err) 55 | } 56 | 57 | s := GetSettings() 58 | if !reflect.DeepEqual(e, s) { 59 | t.Fatalf("Default Settings: want: %+v got: %+v", e, s) 60 | } 61 | } 62 | 63 | func TestSettingsDefault(t *testing.T) { 64 | reset := testSetenv(t, 65 | "USE_STATSD", "", 66 | "STATSD_HOST", "", 67 | "STATSD_PROTOCOL", "", 68 | "STATSD_PORT", "", 69 | "GOSTATS_FLUSH_INTERVAL_SECONDS", "", 70 | "GOSTATS_LOGGING_SINK_DISABLED", "", 71 | ) 72 | defer reset() 73 | exp := Settings{ 74 | UseStatsd: DefaultUseStatsd, 75 | StatsdHost: DefaultStatsdHost, 76 | StatsdProtocol: DefaultStatsdProtocol, 77 | StatsdPort: DefaultStatsdPort, 78 | FlushIntervalS: DefaultFlushIntervalS, 79 | LoggingSinkDisabled: DefaultLoggingSinkDisabled, 80 | } 81 | settings := GetSettings() 82 | if exp != settings { 83 | t.Errorf("Default: want: %+v got: %+v", exp, settings) 84 | } 85 | } 86 | 87 | func TestSettingsOverride(t *testing.T) { 88 | reset := testSetenv(t, 89 | "USE_STATSD", "true", 90 | "STATSD_HOST", "10.0.0.1", 91 | "STATSD_PROTOCOL", "udp", 92 | "STATSD_PORT", "1234", 93 | "GOSTATS_FLUSH_INTERVAL_SECONDS", "3", 94 | "GOSTATS_LOGGING_SINK_DISABLED", "true", 95 | ) 96 | defer reset() 97 | exp := Settings{ 98 | UseStatsd: true, 99 | StatsdHost: "10.0.0.1", 100 | StatsdProtocol: "udp", 101 | StatsdPort: 1234, 102 | FlushIntervalS: 3, 103 | LoggingSinkDisabled: true, 104 | } 105 | settings := GetSettings() 106 | if exp != settings { 107 | t.Errorf("Default: want: %+v got: %+v", exp, settings) 108 | } 109 | } 110 | 111 | func TestSettingsErrors(t *testing.T) { 112 | // STATSD_HOST doesn't error so we don't check it 113 | 114 | tests := map[string]string{ 115 | "USE_STATSD": "FOO!", 116 | "STATSD_PORT": "not-an-int", 117 | "GOSTATS_FLUSH_INTERVAL_SECONDS": "true", 118 | "GOSTATS_LOGGING_SINK_DISABLED": "1337", 119 | } 120 | for key, val := range tests { 121 | t.Run(key, func(t *testing.T) { 122 | reset := testSetenv(t, key, val) 123 | defer reset() 124 | var panicked bool 125 | func() { 126 | defer func() { 127 | panicked = recover() != nil 128 | }() 129 | GetSettings() 130 | }() 131 | if !panicked { 132 | t.Errorf("Settings expected a panic for invalid value %s=%s", key, val) 133 | } 134 | }) 135 | } 136 | } 137 | 138 | func TestFlushInterval(t *testing.T) { 139 | reset := testSetenv(t, 140 | "GOSTATS_FLUSH_INTERVAL_SECONDS", "3", 141 | ) 142 | defer reset() 143 | settings := GetSettings() 144 | interval := settings.FlushInterval() 145 | expected := time.Duration(3) * time.Second 146 | if interval != expected { 147 | t.Errorf("Flush interval does not match expected duration %s != %s", interval, expected) 148 | } 149 | } 150 | -------------------------------------------------------------------------------- /net_util_test.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "bufio" 5 | "bytes" 6 | "errors" 7 | "fmt" 8 | "io" 9 | "net" 10 | "os" 11 | "strconv" 12 | "strings" 13 | "sync" 14 | "syscall" 15 | "testing" 16 | "time" 17 | ) 18 | 19 | type TestConn interface { 20 | Close() (err error) 21 | Address() net.Addr 22 | Reconnect(t testing.TB, s *netTestSink) 23 | Run(t testing.TB) 24 | } 25 | 26 | type udpConn struct { 27 | ll *net.UDPConn 28 | addr *net.UDPAddr 29 | writeStat func(line []byte) 30 | done chan struct{} 31 | } 32 | 33 | func NewUDPConn(t testing.TB, s *netTestSink) TestConn { 34 | l, err := net.ListenUDP("udp", &net.UDPAddr{ 35 | IP: net.IPv4(127, 0, 0, 1), 36 | Port: 0, 37 | }) 38 | if err != nil { 39 | t.Fatal("ListenUDP:", err) 40 | } 41 | c := &udpConn{ 42 | ll: l, 43 | addr: l.LocalAddr().(*net.UDPAddr), 44 | writeStat: s.writeStat, 45 | done: s.done, 46 | } 47 | go c.Run(t) 48 | return c 49 | } 50 | 51 | func (c *udpConn) Address() net.Addr { 52 | return c.ll.LocalAddr() 53 | } 54 | 55 | func (c *udpConn) Close() (err error) { 56 | return c.ll.Close() 57 | } 58 | 59 | func (c *udpConn) Reconnect(t testing.TB, s *netTestSink) { 60 | reconnectRetry(t, func() error { 61 | l, err := net.ListenUDP(c.addr.Network(), c.addr) 62 | if err != nil { 63 | return err 64 | } 65 | c.ll = l 66 | c.writeStat = s.writeStat 67 | c.done = s.done 68 | go c.Run(t) 69 | return nil 70 | }) 71 | } 72 | 73 | func (c *udpConn) Run(t testing.TB) { 74 | defer close(c.done) 75 | buf := bufio.NewReader(c.ll) 76 | var err error 77 | for { 78 | b, e := buf.ReadBytes('\n') 79 | if len(b) > 0 { 80 | c.writeStat(b) 81 | } 82 | if e != nil { 83 | if e != io.EOF { 84 | err = e 85 | } 86 | break 87 | } 88 | } 89 | if err != nil && !strings.Contains(err.Error(), "use of closed network connection") { 90 | t.Logf("Error: reading stats: %v", err) 91 | } 92 | } 93 | 94 | type tcpConn struct { 95 | ll *net.TCPListener 96 | addr *net.TCPAddr 97 | writeStat func(line []byte) 98 | done chan struct{} 99 | } 100 | 101 | func NewTCPConn(t testing.TB, s *netTestSink) TestConn { 102 | l, err := net.ListenTCP("tcp", &net.TCPAddr{ 103 | IP: net.IPv4(127, 0, 0, 1), 104 | Port: 0, 105 | }) 106 | if err != nil { 107 | t.Fatal("ListenTCP:", err) 108 | } 109 | c := &tcpConn{ 110 | ll: l, 111 | addr: l.Addr().(*net.TCPAddr), 112 | writeStat: s.writeStat, 113 | done: s.done, 114 | } 115 | go c.Run(t) 116 | return c 117 | } 118 | 119 | func (c *tcpConn) Address() net.Addr { 120 | return c.ll.Addr() 121 | } 122 | 123 | func (c *tcpConn) Close() (err error) { 124 | return c.ll.Close() 125 | } 126 | 127 | func (c *tcpConn) Reconnect(t testing.TB, s *netTestSink) { 128 | reconnectRetry(t, func() error { 129 | l, err := net.ListenTCP(c.addr.Network(), c.addr) 130 | if err != nil { 131 | return err 132 | } 133 | c.ll = l 134 | c.writeStat = s.writeStat 135 | c.done = s.done 136 | go c.Run(t) 137 | return nil 138 | }) 139 | } 140 | 141 | func (c *tcpConn) Run(t testing.TB) { 142 | defer close(c.done) 143 | buf := bufio.NewReader(nil) 144 | for { 145 | conn, err := c.ll.AcceptTCP() 146 | if err != nil { 147 | // Log errors other than poll.ErrNetClosing, which is an 148 | // internal error so we have to match against it's string. 149 | if !strings.Contains(err.Error(), "use of closed network connection") { 150 | t.Logf("Error: accept: %v", err) 151 | } 152 | return 153 | } 154 | // read stats line by line 155 | buf.Reset(conn) 156 | for { 157 | b, e := buf.ReadBytes('\n') 158 | if len(b) > 0 { 159 | c.writeStat(b) 160 | } 161 | if e != nil { 162 | if e != io.EOF { 163 | err = e 164 | } 165 | break 166 | } 167 | } 168 | if err != nil { 169 | t.Errorf("Error: reading stats: %v", err) 170 | } 171 | } 172 | } 173 | 174 | type netTestSink struct { 175 | conn TestConn 176 | mu sync.Mutex // buf lock 177 | buf bytes.Buffer 178 | stats chan string 179 | done chan struct{} // closed when read loop exits 180 | protocol string 181 | } 182 | 183 | func newNetTestSink(t testing.TB, protocol string) *netTestSink { 184 | s := &netTestSink{ 185 | stats: make(chan string, 64), 186 | done: make(chan struct{}), 187 | protocol: protocol, 188 | } 189 | switch protocol { 190 | case "udp": 191 | s.conn = NewUDPConn(t, s) 192 | case "tcp": 193 | s.conn = NewTCPConn(t, s) 194 | default: 195 | t.Fatalf("invalid network protocol: %q", protocol) 196 | } 197 | return s 198 | } 199 | 200 | func (s *netTestSink) writeStat(line []byte) { 201 | select { 202 | case s.stats <- string(line): 203 | default: 204 | } 205 | s.mu.Lock() 206 | s.buf.Write(line) 207 | s.mu.Unlock() 208 | } 209 | 210 | func (s *netTestSink) Restart(t testing.TB, resetBuffer bool) { 211 | if err := s.Close(); err != nil { 212 | if !strings.Contains(err.Error(), "use of closed network connection") { 213 | t.Fatal(err) 214 | } 215 | } 216 | select { 217 | case <-s.done: 218 | // Ok 219 | case <-time.After(time.Second * 3): 220 | t.Fatal("timeout waiting for run loop to exit") 221 | } 222 | 223 | if resetBuffer { 224 | s.buf.Reset() 225 | } 226 | s.stats = make(chan string, 64) 227 | s.done = make(chan struct{}) 228 | s.conn.Reconnect(t, s) 229 | } 230 | 231 | func (s *netTestSink) WaitForStat(t testing.TB, timeout time.Duration) string { 232 | t.Helper() 233 | if timeout <= 0 { 234 | timeout = defaultRetryInterval * 2 235 | } 236 | to := time.NewTimer(timeout) 237 | defer to.Stop() 238 | select { 239 | case s := <-s.stats: 240 | return s 241 | case <-to.C: 242 | t.Fatalf("timeout waiting to receive stat: %s", timeout) 243 | } 244 | return "" 245 | } 246 | 247 | func (s *netTestSink) Close() error { 248 | select { 249 | case <-s.done: 250 | return nil // closed 251 | default: 252 | return s.conn.Close() // WARN 253 | // return s.ll.Close() // WARN 254 | } 255 | } 256 | 257 | func (s *netTestSink) Stats() <-chan string { 258 | return s.stats 259 | } 260 | 261 | func (s *netTestSink) Bytes() []byte { 262 | s.mu.Lock() 263 | b := append([]byte(nil), s.buf.Bytes()...) 264 | s.mu.Unlock() 265 | return b 266 | } 267 | 268 | func (s *netTestSink) String() string { 269 | s.mu.Lock() 270 | str := s.buf.String() 271 | s.mu.Unlock() 272 | return str 273 | } 274 | 275 | func (s *netTestSink) Host(t testing.TB) string { 276 | t.Helper() 277 | host, _, err := net.SplitHostPort(s.conn.Address().String()) 278 | if err != nil { 279 | t.Fatal(err) 280 | } 281 | return host 282 | } 283 | 284 | func (s *netTestSink) Port(t testing.TB) int { 285 | t.Helper() 286 | _, port, err := net.SplitHostPort(s.conn.Address().String()) 287 | if err != nil { 288 | t.Fatal(err) 289 | } 290 | n, err := strconv.Atoi(port) 291 | if err != nil { 292 | t.Fatal(err) 293 | } 294 | return n 295 | } 296 | 297 | func (s *netTestSink) Protocol() string { 298 | return s.protocol 299 | } 300 | 301 | func mergeEnv(extra ...string) []string { 302 | var prefixes []string 303 | for _, s := range extra { 304 | n := strings.IndexByte(s, '=') 305 | prefixes = append(prefixes, s[:n+1]) 306 | } 307 | ignore := func(s string) bool { 308 | for _, pfx := range prefixes { 309 | if strings.HasPrefix(s, pfx) { 310 | return true 311 | } 312 | } 313 | return false 314 | } 315 | 316 | env := os.Environ() 317 | a := env[:0] 318 | for _, s := range env { 319 | if !ignore(s) { 320 | a = append(a, s) 321 | } 322 | } 323 | return append(a, extra...) 324 | } 325 | 326 | func (s *netTestSink) CommandEnv(t testing.TB) []string { 327 | return mergeEnv( 328 | fmt.Sprintf("STATSD_PORT=%d", s.Port(t)), 329 | fmt.Sprintf("STATSD_HOST=%s", s.Host(t)), 330 | fmt.Sprintf("STATSD_PROTOCOL=%s", s.Protocol()), 331 | "GOSTATS_FLUSH_INTERVAL_SECONDS=1", 332 | ) 333 | } 334 | 335 | func reconnectRetry(t testing.TB, fn func() error) { 336 | const ( 337 | Retry = time.Second / 4 338 | Timeout = 5 * time.Second 339 | N = int(Timeout / Retry) 340 | ) 341 | var err error 342 | for i := 0; i < N; i++ { 343 | err = fn() 344 | if err == nil { 345 | return 346 | } 347 | // Retry if the error is due to the address being in use. 348 | // On slow systems (CI) it can take awhile for the OS to 349 | // realize the address is not in use. 350 | if errors.Is(err, syscall.EADDRINUSE) { 351 | time.Sleep(Retry) 352 | } else { 353 | t.Fatalf("unexpected error reconnecting: %s", err) 354 | return // unreachable 355 | } 356 | } 357 | t.Fatalf("failed to reconnect after %d attempts and %s: %v", 358 | N, Timeout, err) 359 | } 360 | 361 | func TestReconnectRetryTCP(t *testing.T) { 362 | t.Parallel() 363 | 364 | l1, err := net.ListenTCP("tcp", &net.TCPAddr{ 365 | IP: net.IPv4(127, 0, 0, 1), 366 | Port: 0, 367 | }) 368 | if err != nil { 369 | t.Fatal(err) 370 | } 371 | defer l1.Close() 372 | 373 | l2, err := net.ListenTCP(l1.Addr().Network(), l1.Addr().(*net.TCPAddr)) 374 | if err == nil { 375 | l2.Close() 376 | t.Fatal("expected an error got nil") 377 | } 378 | 379 | if !errors.Is(err, syscall.EADDRINUSE) { 380 | t.Fatalf("expected error to wrap %T got: %#v", syscall.EADDRINUSE, err) 381 | } 382 | 383 | first := true 384 | callCount := 0 385 | reconnectRetry(t, func() error { 386 | callCount++ 387 | if first { 388 | first = false 389 | return err 390 | } 391 | return nil 392 | }) 393 | 394 | if callCount != 2 { 395 | t.Errorf("Expected call cound to be %d got: %d", 2, callCount) 396 | } 397 | } 398 | 399 | func TestReconnectRetryUDP(t *testing.T) { 400 | t.Parallel() 401 | 402 | l1, err := net.ListenUDP("udp", &net.UDPAddr{ 403 | IP: net.IPv4(127, 0, 0, 1), 404 | Port: 0, 405 | }) 406 | if err != nil { 407 | t.Fatal(err) 408 | } 409 | defer l1.Close() 410 | 411 | l2, err := net.ListenUDP(l1.LocalAddr().Network(), l1.LocalAddr().(*net.UDPAddr)) 412 | if err == nil { 413 | l2.Close() 414 | t.Fatal("expected an error got nil") 415 | } 416 | 417 | if !errors.Is(err, syscall.EADDRINUSE) { 418 | t.Fatalf("expected error to wrap %T got: %#v", syscall.EADDRINUSE, err) 419 | } 420 | 421 | first := true 422 | callCount := 0 423 | reconnectRetry(t, func() error { 424 | callCount++ 425 | if first { 426 | first = false 427 | return err 428 | } 429 | return nil 430 | }) 431 | 432 | if callCount != 2 { 433 | t.Errorf("Expected call cound to be %d got: %d", 2, callCount) 434 | } 435 | } 436 | -------------------------------------------------------------------------------- /net_sink.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "bufio" 5 | "bytes" 6 | "fmt" 7 | "math" 8 | "net" 9 | "os" 10 | "strconv" 11 | "sync" 12 | "time" 13 | ) 14 | 15 | // Logger is used to log errors and other important operational 16 | // information while using gostats. 17 | // 18 | // For convenience of transitioning from logrus to zap, this interface 19 | // conforms BOTH to logrus.Logger as well as the Zap's Sugared logger. 20 | type Logger interface { 21 | Errorf(msg string, args ...interface{}) 22 | Warnf(msg string, args ...interface{}) 23 | } 24 | 25 | const ( 26 | defaultRetryInterval = time.Second * 3 27 | defaultDialTimeout = defaultRetryInterval / 2 28 | defaultWriteTimeout = time.Second 29 | 30 | flushInterval = time.Second 31 | logOnEveryNDroppedBytes = 1 << 15 // Log once per 32kb of dropped stats 32 | defaultBufferSizeTCP = 1 << 16 33 | 34 | // 1432 bytes is optimal for regular networks with an MTU of 1500 and 35 | // is to prevent fragmenting UDP datagrams 36 | defaultBufferSizeUDP = 1432 37 | 38 | approxMaxMemBytes = 1 << 22 39 | ) 40 | 41 | // An SinkOption configures a Sink. 42 | type SinkOption interface { 43 | apply(*netSink) 44 | } 45 | 46 | // sinkOptionFunc wraps a func so it satisfies the Option interface. 47 | type sinkOptionFunc func(*netSink) 48 | 49 | func (f sinkOptionFunc) apply(sink *netSink) { 50 | f(sink) 51 | } 52 | 53 | // WithStatsdHost sets the host of the statsd sink otherwise the host is 54 | // read from the environment variable "STATSD_HOST". 55 | func WithStatsdHost(host string) SinkOption { 56 | return sinkOptionFunc(func(sink *netSink) { 57 | sink.conf.StatsdHost = host 58 | }) 59 | } 60 | 61 | // WithStatsdProtocol sets the network protocol ("udp" or "tcp") of the statsd 62 | // sink otherwise the protocol is read from the environment variable 63 | // "STATSD_PROTOCOL". 64 | func WithStatsdProtocol(protocol string) SinkOption { 65 | return sinkOptionFunc(func(sink *netSink) { 66 | sink.conf.StatsdProtocol = protocol 67 | }) 68 | } 69 | 70 | // WithStatsdPort sets the port of the statsd sink otherwise the port is 71 | // read from the environment variable "STATSD_PORT". 72 | func WithStatsdPort(port int) SinkOption { 73 | return sinkOptionFunc(func(sink *netSink) { 74 | sink.conf.StatsdPort = port 75 | }) 76 | } 77 | 78 | // WithLogger configures the sink to use the provided logger otherwise 79 | // the built-in zap-like logger is used. 80 | func WithLogger(log Logger) SinkOption { 81 | return sinkOptionFunc(func(sink *netSink) { 82 | sink.log = log 83 | }) 84 | } 85 | 86 | // NewTCPStatsdSink returns a new NetStink. This function name exists for 87 | // backwards compatibility. 88 | func NewTCPStatsdSink(opts ...SinkOption) FlushableSink { 89 | return NewNetSink(opts...) 90 | } 91 | 92 | // NewNetSink returns a FlushableSink that writes to a statsd sink over the 93 | // network. By default settings are taken from the environment, but can be 94 | // overridden via SinkOptions. 95 | func NewNetSink(opts ...SinkOption) FlushableSink { 96 | s := &netSink{ 97 | // arbitrarily buffered 98 | doFlush: make(chan chan struct{}, 8), 99 | 100 | // default logging sink mimics the previously-used logrus 101 | // logger by logging to stderr 102 | log: &loggingSink{writer: os.Stderr, now: time.Now}, 103 | 104 | // TODO (CEV): auto loading from the env is bad and should be removed. 105 | conf: GetSettings(), 106 | } 107 | for _, opt := range opts { 108 | opt.apply(s) 109 | } 110 | 111 | // Calculate buffer size based on protocol, for UDP we want to pick a 112 | // buffer size that will prevent datagram fragmentation. 113 | var bufSize int 114 | switch s.conf.StatsdProtocol { 115 | case "udp", "udp4", "udp6": 116 | bufSize = defaultBufferSizeUDP 117 | default: 118 | bufSize = defaultBufferSizeTCP 119 | } 120 | 121 | s.outc = make(chan *bytes.Buffer, approxMaxMemBytes/bufSize) 122 | s.retryc = make(chan *bytes.Buffer, 1) // It should be okay to limit this given we preferentially process from this over outc. 123 | 124 | writer := &sinkWriter{outc: s.outc} 125 | s.bufWriter = bufio.NewWriterSize(writer, bufSize) 126 | 127 | go s.run() 128 | return s 129 | } 130 | 131 | type netSink struct { 132 | conn net.Conn 133 | outc chan *bytes.Buffer 134 | retryc chan *bytes.Buffer 135 | mu sync.Mutex 136 | bufWriter *bufio.Writer 137 | doFlush chan chan struct{} 138 | droppedBytes uint64 139 | log Logger 140 | conf Settings 141 | } 142 | 143 | type sinkWriter struct { 144 | outc chan<- *bytes.Buffer 145 | } 146 | 147 | func (w *sinkWriter) Write(p []byte) (int, error) { 148 | n := len(p) 149 | dest := getBuffer() 150 | dest.Write(p) 151 | select { 152 | case w.outc <- dest: 153 | return n, nil 154 | default: 155 | return 0, fmt.Errorf("statsd channel full, dropping stats buffer with %d bytes", n) 156 | } 157 | } 158 | 159 | func (s *netSink) Flush() { 160 | if s.flush() != nil { 161 | return // nothing we can do 162 | } 163 | ch := make(chan struct{}) 164 | s.doFlush <- ch 165 | <-ch 166 | } 167 | 168 | func (s *netSink) flush() error { 169 | s.mu.Lock() 170 | err := s.bufWriter.Flush() 171 | if err != nil { 172 | s.handleFlushError(err) 173 | } 174 | s.mu.Unlock() 175 | return err 176 | } 177 | 178 | func (s *netSink) drainFlushQueue() { 179 | // Limit the number of items we'll flush to prevent this from possibly 180 | // hanging when the flush channel is saturated with sends. 181 | doFlush := s.doFlush 182 | n := cap(doFlush) * 8 183 | for i := 0; i < n; i++ { 184 | select { 185 | case ch := <-doFlush: 186 | close(ch) 187 | default: 188 | return 189 | } 190 | } 191 | } 192 | 193 | // s.mu should be held 194 | func (s *netSink) handleFlushErrorSize(err error, dropped int) { 195 | d := uint64(dropped) 196 | if (s.droppedBytes+d)%logOnEveryNDroppedBytes > s.droppedBytes%logOnEveryNDroppedBytes { 197 | s.log.Errorf("dropped %d bytes: %s", s.droppedBytes+d, err) 198 | } 199 | s.droppedBytes += d 200 | 201 | s.bufWriter.Reset(&sinkWriter{ 202 | outc: s.outc, 203 | }) 204 | } 205 | 206 | // s.mu should be held 207 | func (s *netSink) handleFlushError(err error) { 208 | s.handleFlushErrorSize(err, s.bufWriter.Buffered()) 209 | } 210 | 211 | func (s *netSink) writeBuffer(b *buffer) { 212 | s.mu.Lock() 213 | if s.bufWriter.Available() < b.Len() { 214 | if err := s.bufWriter.Flush(); err != nil { 215 | s.handleFlushError(err) 216 | } 217 | // If there is an error we reset the bufWriter so its 218 | // okay to attempt the write after the failed flush. 219 | } 220 | if _, err := s.bufWriter.Write(*b); err != nil { 221 | s.handleFlushError(err) 222 | } 223 | s.mu.Unlock() 224 | } 225 | 226 | func (s *netSink) flushUint64(name, suffix string, u uint64) { 227 | b := pbFree.Get().(*buffer) 228 | 229 | b.WriteString(name) 230 | b.WriteChar(':') 231 | b.WriteUnit64(u) 232 | b.WriteString(suffix) 233 | 234 | s.writeBuffer(b) 235 | 236 | b.Reset() 237 | pbFree.Put(b) 238 | } 239 | 240 | func (s *netSink) flushFloat64(name, suffix string, f float64) { 241 | b := pbFree.Get().(*buffer) 242 | 243 | b.WriteString(name) 244 | b.WriteChar(':') 245 | b.WriteFloat64(f) 246 | b.WriteString(suffix) 247 | 248 | s.writeBuffer(b) 249 | 250 | b.Reset() 251 | pbFree.Put(b) 252 | } 253 | 254 | func (s *netSink) FlushCounter(name string, value uint64) { 255 | s.flushUint64(name, "|c\n", value) 256 | } 257 | 258 | func (s *netSink) FlushGauge(name string, value uint64) { 259 | s.flushUint64(name, "|g\n", value) 260 | } 261 | 262 | func (s *netSink) FlushTimer(name string, value float64) { 263 | // Since we mistakenly use floating point values to represent time 264 | // durations this method is often passed an integer encoded as a 265 | // float. Formatting integers is much faster (>2x) than formatting 266 | // floats so use integer formatting whenever possible. 267 | // 268 | if 0 <= value && value < math.MaxUint64 && math.Trunc(value) == value { 269 | s.flushUint64(name, "|ms\n", uint64(value)) 270 | } else { 271 | s.flushFloat64(name, "|ms\n", value) 272 | } 273 | } 274 | 275 | func (s *netSink) run() { 276 | addr := net.JoinHostPort(s.conf.StatsdHost, strconv.Itoa(s.conf.StatsdPort)) 277 | 278 | var reconnectFailed bool // true if last reconnect failed 279 | 280 | t := time.NewTicker(flushInterval) 281 | defer t.Stop() 282 | for { 283 | if s.conn == nil { 284 | if err := s.connect(addr); err != nil { 285 | s.log.Warnf("connection error: %s", err) 286 | 287 | // If the previous reconnect attempt failed, drain the flush 288 | // queue to prevent Flush() from blocking indefinitely. 289 | if reconnectFailed { 290 | s.drainFlushQueue() 291 | } 292 | reconnectFailed = true 293 | 294 | // TODO (CEV): don't sleep on the first retry 295 | time.Sleep(defaultRetryInterval) 296 | continue 297 | } 298 | reconnectFailed = false 299 | } 300 | 301 | // Handle buffers that need to be retried first, if they exist. 302 | select { 303 | case buf := <-s.retryc: 304 | if err := s.writeToConn(buf); err != nil { 305 | s.mu.Lock() 306 | s.handleFlushErrorSize(err, buf.Len()) 307 | s.mu.Unlock() 308 | } 309 | putBuffer(buf) 310 | continue 311 | default: 312 | // Drop through in case retryc has nothing. 313 | } 314 | 315 | select { 316 | case <-t.C: 317 | s.flush() 318 | case done := <-s.doFlush: 319 | // Only flush pending buffers, this prevents an issue where 320 | // continuous writes prevent the flush loop from exiting. 321 | // 322 | // If there is an error writeToConn() will set the conn to 323 | // nil thus breaking the loop. 324 | // 325 | n := len(s.outc) 326 | for i := 0; i < n && s.conn != nil; i++ { 327 | buf := <-s.outc 328 | if err := s.writeToConn(buf); err != nil { 329 | s.retryc <- buf 330 | continue 331 | } 332 | putBuffer(buf) 333 | } 334 | close(done) 335 | case buf := <-s.outc: 336 | if err := s.writeToConn(buf); err != nil { 337 | s.retryc <- buf 338 | continue 339 | } 340 | putBuffer(buf) 341 | } 342 | } 343 | } 344 | 345 | // writeToConn writes the buffer to the underlying conn. May only be called 346 | // from run(). 347 | func (s *netSink) writeToConn(buf *bytes.Buffer) error { 348 | // TODO (CEV): parameterize timeout 349 | s.conn.SetWriteDeadline(time.Now().Add(defaultWriteTimeout)) 350 | _, err := buf.WriteTo(s.conn) 351 | s.conn.SetWriteDeadline(time.Time{}) // clear 352 | 353 | if err != nil { 354 | _ = s.conn.Close() 355 | s.conn = nil // this will break the loop 356 | } 357 | return err 358 | } 359 | 360 | func (s *netSink) connect(address string) error { 361 | // TODO (CEV): parameterize timeout 362 | conn, err := net.DialTimeout(s.conf.StatsdProtocol, address, defaultDialTimeout) 363 | if err == nil { 364 | s.conn = conn 365 | } 366 | return err 367 | } 368 | 369 | var bufferPool sync.Pool 370 | 371 | func getBuffer() *bytes.Buffer { 372 | if v := bufferPool.Get(); v != nil { 373 | b := v.(*bytes.Buffer) 374 | b.Reset() 375 | return b 376 | } 377 | return new(bytes.Buffer) 378 | } 379 | 380 | func putBuffer(b *bytes.Buffer) { 381 | bufferPool.Put(b) 382 | } 383 | 384 | // pbFree is the print buffer pool 385 | var pbFree = sync.Pool{ 386 | New: func() interface{} { 387 | b := make(buffer, 0, 128) 388 | return &b 389 | }, 390 | } 391 | 392 | // Use a fast and simple buffer for constructing statsd messages 393 | type buffer []byte 394 | 395 | func (b *buffer) Len() int { return len(*b) } 396 | 397 | func (b *buffer) Reset() { *b = (*b)[:0] } 398 | 399 | func (b *buffer) Write(p []byte) { 400 | *b = append(*b, p...) 401 | } 402 | 403 | func (b *buffer) WriteString(s string) { 404 | *b = append(*b, s...) 405 | } 406 | 407 | // This is named WriteChar instead of WriteByte because the 'stdmethods' check 408 | // of 'go vet' wants WriteByte to have the signature: 409 | // 410 | // func (b *buffer) WriteByte(c byte) error { ... } 411 | func (b *buffer) WriteChar(c byte) { 412 | *b = append(*b, c) 413 | } 414 | 415 | func (b *buffer) WriteUnit64(val uint64) { 416 | *b = strconv.AppendUint(*b, val, 10) 417 | } 418 | 419 | func (b *buffer) WriteFloat64(val float64) { 420 | *b = strconv.AppendFloat(*b, val, 'f', 6, 64) 421 | } 422 | -------------------------------------------------------------------------------- /mock/sink_test.go: -------------------------------------------------------------------------------- 1 | package mock_test 2 | 3 | import ( 4 | "fmt" 5 | "math" 6 | "reflect" 7 | "runtime" 8 | "sync" 9 | "sync/atomic" 10 | "testing" 11 | 12 | "github.com/lyft/gostats/mock" 13 | ) 14 | 15 | type ErrorTest struct { 16 | testing.TB 17 | errMsg string 18 | } 19 | 20 | func NewErrorTest(t testing.TB) *ErrorTest { 21 | return &ErrorTest{TB: t} 22 | } 23 | 24 | func (t *ErrorTest) Error(args ...interface{}) { 25 | t.errMsg = fmt.Sprint(args...) 26 | } 27 | 28 | func (t *ErrorTest) Errorf(format string, args ...interface{}) { 29 | t.errMsg = fmt.Sprintf(format, args...) 30 | } 31 | 32 | func (t *ErrorTest) AssertErrorMsg(format string, args ...interface{}) { 33 | exp := fmt.Sprintf(format, args...) 34 | if t.errMsg != exp { 35 | t.TB.Errorf("Expected error message: `%s` got: `%s`", exp, t.errMsg) 36 | } 37 | } 38 | 39 | func AssertErrorMsg(t testing.TB, fn func(t testing.TB), format string, args ...interface{}) { 40 | t.Helper() 41 | x := NewErrorTest(t) 42 | fn(x) 43 | x.AssertErrorMsg(format, args...) 44 | } 45 | 46 | func (t *ErrorTest) Reset() testing.TB { 47 | t.errMsg = "" 48 | return t 49 | } 50 | 51 | func TestSink(t *testing.T) { 52 | testCounter := func(t *testing.T, exp uint64, sink *mock.Sink) { 53 | t.Helper() 54 | const name = "test-counter" 55 | sink.FlushCounter(name, exp) 56 | sink.AssertCounterExists(t, name) 57 | sink.AssertCounterEquals(t, name, exp) 58 | sink.AssertCounterCallCount(t, name, 1) 59 | if n := sink.Counter(name); n != exp { 60 | t.Errorf("Counter(): want: %d got: %d", exp, n) 61 | } 62 | 63 | const missing = name + "-MISSING" 64 | sink.AssertCounterNotExists(t, missing) 65 | 66 | fns := []func(t testing.TB){ 67 | func(t testing.TB) { sink.AssertCounterExists(t, missing) }, 68 | func(t testing.TB) { sink.AssertCounterEquals(t, missing, 9999) }, 69 | func(t testing.TB) { sink.AssertCounterCallCount(t, missing, 9999) }, 70 | } 71 | for _, fn := range fns { 72 | AssertErrorMsg(t, fn, "gostats/mock: Counter (%q): not found in: [\"test-counter\"]", missing) 73 | } 74 | 75 | AssertErrorMsg(t, func(t testing.TB) { 76 | sink.AssertCounterEquals(t, name, 9999) 77 | }, "gostats/mock: Counter (%q): Expected: %d Got: %d", name, 9999, exp) 78 | 79 | AssertErrorMsg(t, func(t testing.TB) { 80 | sink.AssertCounterCallCount(t, name, 9999) 81 | }, "gostats/mock: Counter (%q) Call Count: Expected: %d Got: %d", name, 9999, 1) 82 | 83 | AssertErrorMsg(t, func(t testing.TB) { 84 | sink.AssertCounterNotExists(t, name) 85 | }, "gostats/mock: Counter (%q): expected Counter to not exist", name) 86 | } 87 | 88 | testGauge := func(t *testing.T, exp uint64, sink *mock.Sink) { 89 | const name = "test-gauge" 90 | sink.FlushGauge(name, exp) 91 | sink.AssertGaugeExists(t, name) 92 | sink.AssertGaugeEquals(t, name, exp) 93 | sink.AssertGaugeCallCount(t, name, 1) 94 | if n := sink.Gauge(name); n != exp { 95 | t.Errorf("Gauge(): want: %d got: %d", exp, n) 96 | } 97 | 98 | const missing = name + "-MISSING" 99 | sink.AssertGaugeNotExists(t, missing) 100 | 101 | fns := []func(t testing.TB){ 102 | func(t testing.TB) { sink.AssertGaugeExists(t, missing) }, 103 | func(t testing.TB) { sink.AssertGaugeEquals(t, missing, 9999) }, 104 | func(t testing.TB) { sink.AssertGaugeCallCount(t, missing, 9999) }, 105 | } 106 | for _, fn := range fns { 107 | AssertErrorMsg(t, fn, "gostats/mock: Gauge (%q): not found in: [\"test-gauge\"]", missing) 108 | } 109 | 110 | AssertErrorMsg(t, func(t testing.TB) { 111 | sink.AssertGaugeEquals(t, name, 9999) 112 | }, "gostats/mock: Gauge (%q): Expected: %d Got: %d", name, 9999, exp) 113 | 114 | AssertErrorMsg(t, func(t testing.TB) { 115 | sink.AssertGaugeCallCount(t, name, 9999) 116 | }, "gostats/mock: Gauge (%q) Call Count: Expected: %d Got: %d", name, 9999, 1) 117 | 118 | AssertErrorMsg(t, func(t testing.TB) { 119 | sink.AssertGaugeNotExists(t, name) 120 | }, "gostats/mock: Gauge (%q): expected Gauge to not exist", name) 121 | } 122 | 123 | testTimer := func(t *testing.T, exp float64, sink *mock.Sink) { 124 | const name = "test-timer" 125 | sink.FlushTimer(name, exp) 126 | sink.AssertTimerExists(t, name) 127 | sink.AssertTimerEquals(t, name, exp) 128 | sink.AssertTimerCallCount(t, name, 1) 129 | if n := sink.Timer(name); n != exp { 130 | t.Errorf("Timer(): want: %f got: %f", exp, n) 131 | } 132 | 133 | const missing = name + "-MISSING" 134 | sink.AssertTimerNotExists(t, missing) 135 | 136 | fns := []func(t testing.TB){ 137 | func(t testing.TB) { sink.AssertTimerExists(t, missing) }, 138 | func(t testing.TB) { sink.AssertTimerEquals(t, missing, 9999) }, 139 | func(t testing.TB) { sink.AssertTimerCallCount(t, missing, 9999) }, 140 | } 141 | for _, fn := range fns { 142 | AssertErrorMsg(t, fn, "gostats/mock: Timer (%q): not found in: [\"test-timer\"]", missing) 143 | } 144 | 145 | AssertErrorMsg(t, func(t testing.TB) { 146 | sink.AssertTimerEquals(t, name, 9999) 147 | }, "gostats/mock: Timer (%q): Expected: %f Got: %f", name, 9999.0, exp) 148 | 149 | AssertErrorMsg(t, func(t testing.TB) { 150 | sink.AssertTimerCallCount(t, name, 9999) 151 | }, "gostats/mock: Timer (%q) Call Count: Expected: %d Got: %d", name, 9999, 1) 152 | 153 | AssertErrorMsg(t, func(t testing.TB) { 154 | sink.AssertTimerNotExists(t, name) 155 | }, "gostats/mock: Timer (%q): expected Timer to not exist", name) 156 | } 157 | 158 | // test 0..1 - we want to make sure that 0 still registers a stat 159 | for i := 0; i < 2; i++ { 160 | t.Run("Counter", func(t *testing.T) { 161 | testCounter(t, uint64(i), mock.NewSink()) 162 | }) 163 | t.Run("Gauge", func(t *testing.T) { 164 | testGauge(t, uint64(i), mock.NewSink()) 165 | }) 166 | t.Run("Timer", func(t *testing.T) { 167 | testTimer(t, float64(i), mock.NewSink()) 168 | }) 169 | // all together now 170 | sink := mock.NewSink() 171 | testCounter(t, 1, sink) 172 | testGauge(t, 1, sink) 173 | testTimer(t, 1, sink) 174 | } 175 | } 176 | 177 | func TestSinkMap_Values(t *testing.T) { 178 | expCounters := make(map[string]uint64) 179 | expGauges := make(map[string]uint64) 180 | expTimers := make(map[string]float64) 181 | 182 | sink := mock.NewSink() 183 | for i := 0; i < 2; i++ { 184 | expCounters[fmt.Sprintf("counter-%d", i)] = uint64(i) 185 | expGauges[fmt.Sprintf("gauge-%d", i)] = uint64(i) 186 | expTimers[fmt.Sprintf("timer-%d", i)] = float64(i) 187 | 188 | sink.FlushCounter(fmt.Sprintf("counter-%d", i), uint64(i)) 189 | sink.FlushGauge(fmt.Sprintf("gauge-%d", i), uint64(i)) 190 | sink.FlushTimer(fmt.Sprintf("timer-%d", i), float64(i)) 191 | } 192 | counters := sink.Counters() 193 | if !reflect.DeepEqual(expCounters, counters) { 194 | t.Errorf("Counters: want: %#v got: %#v", expCounters, counters) 195 | } 196 | 197 | gauges := sink.Gauges() 198 | if !reflect.DeepEqual(expGauges, gauges) { 199 | t.Errorf("Gauges: want: %#v got: %#v", expGauges, gauges) 200 | } 201 | 202 | timers := sink.Timers() 203 | if !reflect.DeepEqual(expTimers, timers) { 204 | t.Errorf("Timers: want: %#v got: %#v", expTimers, timers) 205 | } 206 | 207 | sink.Reset() 208 | if n := len(sink.Counters()); n != 0 { 209 | t.Errorf("Failed to reset Counters got: %d", n) 210 | } 211 | if n := len(sink.Gauges()); n != 0 { 212 | t.Errorf("Failed to reset Gauges got: %d", n) 213 | } 214 | if n := len(sink.Timers()); n != 0 { 215 | t.Errorf("Failed to reset Timers got: %d", n) 216 | } 217 | } 218 | 219 | // Test that the zero Sink is ready for use. 220 | func TestSinkLazyInit(t *testing.T) { 221 | var s mock.Sink 222 | s.FlushCounter("counter", 1) 223 | s.AssertCounterEquals(t, "counter", 1) 224 | } 225 | 226 | func TestFlushTimer(t *testing.T) { 227 | sink := mock.NewSink() 228 | var exp float64 229 | for i := 0; i < 10000; i++ { 230 | sink.FlushTimer("timer", 1) 231 | exp++ 232 | } 233 | sink.AssertTimerEquals(t, "timer", exp) 234 | 235 | // test limits 236 | 237 | sink.Reset() 238 | sink.FlushTimer("timer", math.MaxFloat64) 239 | sink.AssertTimerEquals(t, "timer", math.MaxFloat64) 240 | 241 | sink.Reset() 242 | sink.FlushTimer("timer", math.SmallestNonzeroFloat64) 243 | sink.AssertTimerEquals(t, "timer", math.SmallestNonzeroFloat64) 244 | } 245 | 246 | func TestSink_ThreadSafe(t *testing.T) { 247 | const N = 2000 248 | sink := mock.NewSink() 249 | var ( 250 | counterCalls = new(int64) 251 | gaugeCalls = new(int64) 252 | timerCalls = new(int64) 253 | counterVal = new(uint64) 254 | ) 255 | funcs := [...]func(){ 256 | func() { 257 | atomic.AddInt64(counterCalls, 1) 258 | atomic.AddUint64(counterVal, 1) 259 | sink.FlushCounter("name", 1) 260 | }, 261 | func() { 262 | atomic.AddInt64(gaugeCalls, 1) 263 | sink.FlushGauge("name", 1) 264 | }, 265 | func() { 266 | atomic.AddInt64(timerCalls, 1) 267 | sink.FlushTimer("name", 1) 268 | }, 269 | } 270 | numCPU := runtime.NumCPU() 271 | if numCPU < 2 { 272 | numCPU = 2 273 | } 274 | var wg sync.WaitGroup 275 | for i := 0; i < numCPU; i++ { 276 | wg.Add(1) 277 | go func() { 278 | defer wg.Done() 279 | for i := 0; i < N; i++ { 280 | funcs[i%len(funcs)]() 281 | } 282 | }() 283 | } 284 | wg.Wait() 285 | 286 | sink.AssertCounterCallCount(t, "name", int(atomic.LoadInt64(counterCalls))) 287 | sink.AssertGaugeCallCount(t, "name", int(atomic.LoadInt64(gaugeCalls))) 288 | sink.AssertTimerCallCount(t, "name", int(atomic.LoadInt64(timerCalls))) 289 | 290 | sink.AssertCounterEquals(t, "name", atomic.LoadUint64(counterVal)) 291 | sink.AssertGaugeEquals(t, "name", uint64(1)) 292 | } 293 | 294 | func TestSink_ThreadSafe_Reset(_ *testing.T) { 295 | const N = 2000 296 | sink := mock.NewSink() 297 | funcs := [...]func(){ 298 | func() { sink.Flush() }, 299 | func() { sink.FlushCounter("name", 1) }, 300 | func() { sink.FlushGauge("name", 1) }, 301 | func() { sink.FlushTimer("name", 1) }, 302 | func() { sink.LoadCounter("name") }, 303 | func() { sink.LoadGauge("name") }, 304 | func() { sink.LoadTimer("name") }, 305 | func() { sink.Counter("name") }, 306 | func() { sink.Gauge("name") }, 307 | func() { sink.Timer("name") }, 308 | func() { sink.CounterCallCount("name") }, 309 | func() { sink.GaugeCallCount("name") }, 310 | func() { sink.TimerCallCount("name") }, 311 | } 312 | numCPU := runtime.NumCPU() - 1 313 | if numCPU < 2 { 314 | numCPU = 2 315 | } 316 | done := make(chan struct{}) 317 | go func() { 318 | for { 319 | select { 320 | case <-done: 321 | return 322 | default: 323 | sink.Reset() 324 | } 325 | } 326 | }() 327 | var wg sync.WaitGroup 328 | for i := 0; i < numCPU; i++ { 329 | wg.Add(1) 330 | go func() { 331 | defer wg.Done() 332 | for i := 0; i < N; i++ { 333 | funcs[i%len(funcs)]() 334 | } 335 | }() 336 | } 337 | wg.Wait() 338 | close(done) 339 | } 340 | 341 | // TestFatalExample is an example usage of Fatal() 342 | func TestFatalExample(t *testing.T) { 343 | sink := mock.NewSink() 344 | sink.FlushCounter("name", 1) 345 | sink.AssertCounterEquals(mock.Fatal(t), "name", 1) 346 | } 347 | 348 | func setupBenchmark(prefix string) (*mock.Sink, [128]string) { 349 | var names [128]string 350 | if prefix == "" { 351 | prefix = "mock_sink" 352 | } 353 | for i := 0; i < len(names); i++ { 354 | names[i] = fmt.Sprintf("%s_%d", prefix, i) 355 | } 356 | sink := mock.NewSink() 357 | return sink, names 358 | } 359 | 360 | func BenchmarkFlushCounter(b *testing.B) { 361 | sink, names := setupBenchmark("counter") 362 | b.ResetTimer() 363 | for i := 0; i < b.N; i++ { 364 | sink.FlushCounter(names[i%len(names)], uint64(i)) 365 | } 366 | } 367 | 368 | func BenchmarkFlushCounter_Parallel(b *testing.B) { 369 | sink, names := setupBenchmark("counter") 370 | b.ResetTimer() 371 | b.RunParallel(func(pb *testing.PB) { 372 | for i := 0; pb.Next(); i++ { 373 | sink.FlushCounter(names[i%len(names)], uint64(i)) 374 | } 375 | }) 376 | } 377 | 378 | func BenchmarkFlushTimer(b *testing.B) { 379 | const f = 1234.5678 380 | sink, names := setupBenchmark("timer") 381 | b.ResetTimer() 382 | for i := 0; i < b.N; i++ { 383 | sink.FlushTimer(names[i%len(names)], f) 384 | } 385 | } 386 | 387 | func BenchmarkFlushTimer_Parallel(b *testing.B) { 388 | const f = 1234.5678 389 | sink, names := setupBenchmark("timer") 390 | b.ResetTimer() 391 | b.RunParallel(func(pb *testing.PB) { 392 | for i := 0; pb.Next(); i++ { 393 | sink.FlushTimer(names[i%len(names)], f) 394 | } 395 | }) 396 | } 397 | -------------------------------------------------------------------------------- /mock/sink.go: -------------------------------------------------------------------------------- 1 | package mock 2 | 3 | import ( 4 | "fmt" 5 | "math" 6 | "sync" 7 | "sync/atomic" 8 | "testing" 9 | 10 | "github.com/lyft/gostats/internal/tags" 11 | ) 12 | 13 | type entry struct { 14 | val uint64 15 | count int64 16 | } 17 | 18 | type sink struct { 19 | counters sync.Map 20 | timers sync.Map 21 | gauges sync.Map 22 | } 23 | 24 | // A Sink is a mock sink meant for testing that is safe for concurrent use. 25 | type Sink struct { 26 | store atomic.Value 27 | once sync.Once 28 | } 29 | 30 | func (s *Sink) sink() *sink { 31 | s.once.Do(func() { s.store.Store(new(sink)) }) 32 | return s.store.Load().(*sink) 33 | } 34 | 35 | func (s *Sink) counters() *sync.Map { return &s.sink().counters } 36 | func (s *Sink) timers() *sync.Map { return &s.sink().timers } 37 | func (s *Sink) gauges() *sync.Map { return &s.sink().gauges } 38 | 39 | // NewSink returns a new Sink which implements the stats.Sink interface and is 40 | // suitable for testing. 41 | func NewSink() *Sink { 42 | s := &Sink{} 43 | s.sink() // lazy init 44 | return s 45 | } 46 | 47 | // Flush is a no-op method 48 | func (*Sink) Flush() {} 49 | 50 | // Reset resets the Sink's counters, timers and gauges to zero. 51 | func (s *Sink) Reset() { 52 | s.store.Store(new(sink)) 53 | } 54 | 55 | // FlushCounter implements the stats.Sink.FlushCounter method and adds val to 56 | // stat name. 57 | func (s *Sink) FlushCounter(name string, val uint64) { 58 | counters := s.counters() 59 | v, ok := counters.Load(name) 60 | if !ok { 61 | v, _ = counters.LoadOrStore(name, new(entry)) 62 | } 63 | p := v.(*entry) 64 | atomic.AddUint64(&p.val, val) 65 | atomic.AddInt64(&p.count, 1) 66 | } 67 | 68 | // FlushGauge implements the stats.Sink.FlushGauge method and adds val to 69 | // stat name. 70 | func (s *Sink) FlushGauge(name string, val uint64) { 71 | gauges := s.gauges() 72 | v, ok := gauges.Load(name) 73 | if !ok { 74 | v, _ = gauges.LoadOrStore(name, new(entry)) 75 | } 76 | p := v.(*entry) 77 | atomic.StoreUint64(&p.val, val) 78 | atomic.AddInt64(&p.count, 1) 79 | } 80 | 81 | func atomicAddFloat64(dest *uint64, delta float64) { 82 | for { 83 | cur := atomic.LoadUint64(dest) 84 | curVal := math.Float64frombits(cur) 85 | nxtVal := curVal + delta 86 | nxt := math.Float64bits(nxtVal) 87 | if atomic.CompareAndSwapUint64(dest, cur, nxt) { 88 | return 89 | } 90 | } 91 | } 92 | 93 | // FlushTimer implements the stats.Sink.FlushTimer method and adds val to 94 | // stat name. 95 | func (s *Sink) FlushTimer(name string, val float64) { 96 | timers := s.timers() 97 | v, ok := timers.Load(name) 98 | if !ok { 99 | v, _ = timers.LoadOrStore(name, new(entry)) 100 | } 101 | p := v.(*entry) 102 | atomicAddFloat64(&p.val, val) 103 | atomic.AddInt64(&p.count, 1) 104 | } 105 | 106 | // LoadCounter returns the value for stat name and if it was found. 107 | func (s *Sink) LoadCounter(name string) (uint64, bool) { 108 | v, ok := s.counters().Load(name) 109 | if ok { 110 | p := v.(*entry) 111 | return atomic.LoadUint64(&p.val), true 112 | } 113 | return 0, false 114 | } 115 | 116 | // LoadGauge returns the value for stat name and if it was found. 117 | func (s *Sink) LoadGauge(name string) (uint64, bool) { 118 | v, ok := s.gauges().Load(name) 119 | if ok { 120 | p := v.(*entry) 121 | return atomic.LoadUint64(&p.val), true 122 | } 123 | return 0, false 124 | } 125 | 126 | // LoadTimer returns the value for stat name and if it was found. 127 | func (s *Sink) LoadTimer(name string) (float64, bool) { 128 | v, ok := s.timers().Load(name) 129 | if ok { 130 | p := v.(*entry) 131 | bits := atomic.LoadUint64(&p.val) 132 | return math.Float64frombits(bits), true 133 | } 134 | return 0, false 135 | } 136 | 137 | // ListCounters returns a list of existing counter names. 138 | func (s *Sink) ListCounters() []string { 139 | return keys(s.counters()) 140 | } 141 | 142 | // ListGauges returns a list of existing gauge names. 143 | func (s *Sink) ListGauges() []string { 144 | return keys(s.gauges()) 145 | } 146 | 147 | // ListTimers returns a list of existing timer names. 148 | func (s *Sink) ListTimers() []string { 149 | return keys(s.timers()) 150 | } 151 | 152 | // Note, this may return an incoherent snapshot if contents is being concurrently modified 153 | func keys(m *sync.Map) (a []string) { 154 | m.Range(func(key interface{}, _ interface{}) bool { 155 | a = append(a, key.(string)) 156 | return true 157 | }) 158 | return a 159 | } 160 | 161 | // Counters returns all the counters currently stored by the sink. 162 | func (s *Sink) Counters() map[string]uint64 { 163 | m := make(map[string]uint64) 164 | s.counters().Range(func(k, v interface{}) bool { 165 | p := v.(*entry) 166 | m[k.(string)] = atomic.LoadUint64(&p.val) 167 | return true 168 | }) 169 | return m 170 | } 171 | 172 | // Gauges returns all the gauges currently stored by the sink. 173 | func (s *Sink) Gauges() map[string]uint64 { 174 | m := make(map[string]uint64) 175 | s.gauges().Range(func(k, v interface{}) bool { 176 | p := v.(*entry) 177 | m[k.(string)] = atomic.LoadUint64(&p.val) 178 | return true 179 | }) 180 | return m 181 | } 182 | 183 | // Timers returns all the timers currently stored by the sink. 184 | func (s *Sink) Timers() map[string]float64 { 185 | m := make(map[string]float64) 186 | s.timers().Range(func(k, v interface{}) bool { 187 | p := v.(*entry) 188 | bits := atomic.LoadUint64(&p.val) 189 | m[k.(string)] = math.Float64frombits(bits) 190 | return true 191 | }) 192 | return m 193 | } 194 | 195 | // short-hand methods 196 | 197 | // Counter is shorthand for LoadCounter, zero is returned if the stat is not found. 198 | func (s *Sink) Counter(name string) uint64 { 199 | v, _ := s.LoadCounter(name) 200 | return v 201 | } 202 | 203 | // Gauge is shorthand for LoadGauge, zero is returned if the stat is not found. 204 | func (s *Sink) Gauge(name string) uint64 { 205 | v, _ := s.LoadGauge(name) 206 | return v 207 | } 208 | 209 | // Timer is shorthand for LoadTimer, zero is returned if the stat is not found. 210 | func (s *Sink) Timer(name string) float64 { 211 | v, _ := s.LoadTimer(name) 212 | return v 213 | } 214 | 215 | // these methods are mostly useful for testing 216 | 217 | // CounterCallCount returns the number of times stat name has been called/updated. 218 | func (s *Sink) CounterCallCount(name string) int64 { 219 | v, ok := s.counters().Load(name) 220 | if ok { 221 | return atomic.LoadInt64(&v.(*entry).count) 222 | } 223 | return 0 224 | } 225 | 226 | // GaugeCallCount returns the number of times stat name has been called/updated. 227 | func (s *Sink) GaugeCallCount(name string) int64 { 228 | v, ok := s.gauges().Load(name) 229 | if ok { 230 | return atomic.LoadInt64(&v.(*entry).count) 231 | } 232 | return 0 233 | } 234 | 235 | // TimerCallCount returns the number of times stat name has been called/updated. 236 | func (s *Sink) TimerCallCount(name string) int64 { 237 | v, ok := s.timers().Load(name) 238 | if ok { 239 | return atomic.LoadInt64(&v.(*entry).count) 240 | } 241 | return 0 242 | } 243 | 244 | // test helpers 245 | 246 | // AssertCounterEquals asserts that Counter name is present and has value exp. 247 | func (s *Sink) AssertCounterEquals(tb testing.TB, name string, exp uint64) { 248 | tb.Helper() 249 | u, ok := s.LoadCounter(name) 250 | if !ok { 251 | tb.Errorf("gostats/mock: Counter (%q): not found in: %q", name, s.ListCounters()) 252 | return 253 | } 254 | if u != exp { 255 | tb.Errorf("gostats/mock: Counter (%q): Expected: %d Got: %d", name, exp, u) 256 | } 257 | } 258 | 259 | // AssertGaugeEquals asserts that Gauge name is present and has value exp. 260 | func (s *Sink) AssertGaugeEquals(tb testing.TB, name string, exp uint64) { 261 | tb.Helper() 262 | u, ok := s.LoadGauge(name) 263 | if !ok { 264 | tb.Errorf("gostats/mock: Gauge (%q): not found in: %q", name, s.ListGauges()) 265 | return 266 | } 267 | if u != exp { 268 | tb.Errorf("gostats/mock: Gauge (%q): Expected: %d Got: %d", name, exp, u) 269 | } 270 | } 271 | 272 | // AssertTimerEquals asserts that Timer name is present and has value exp. 273 | func (s *Sink) AssertTimerEquals(tb testing.TB, name string, exp float64) { 274 | tb.Helper() 275 | f, ok := s.LoadTimer(name) 276 | if !ok { 277 | tb.Errorf("gostats/mock: Timer (%q): not found in: %q", name, s.ListTimers()) 278 | return 279 | } 280 | if f != exp { 281 | tb.Errorf("gostats/mock: Timer (%q): Expected: %f Got: %f", name, exp, f) 282 | } 283 | } 284 | 285 | // AssertCounterExists asserts that Counter name exists. 286 | func (s *Sink) AssertCounterExists(tb testing.TB, name string) { 287 | tb.Helper() 288 | if _, ok := s.LoadCounter(name); !ok { 289 | tb.Errorf("gostats/mock: Counter (%q): not found in: %q", name, s.ListCounters()) 290 | } 291 | } 292 | 293 | // AssertGaugeExists asserts that Gauge name exists. 294 | func (s *Sink) AssertGaugeExists(tb testing.TB, name string) { 295 | tb.Helper() 296 | if _, ok := s.LoadGauge(name); !ok { 297 | tb.Errorf("gostats/mock: Gauge (%q): not found in: %q", name, s.ListGauges()) 298 | } 299 | } 300 | 301 | // AssertTimerExists asserts that Timer name exists. 302 | func (s *Sink) AssertTimerExists(tb testing.TB, name string) { 303 | tb.Helper() 304 | if _, ok := s.LoadTimer(name); !ok { 305 | tb.Errorf("gostats/mock: Timer (%q): not found in: %q", name, s.ListTimers()) 306 | } 307 | } 308 | 309 | // AssertCounterNotExists asserts that Counter name does not exist. 310 | func (s *Sink) AssertCounterNotExists(tb testing.TB, name string) { 311 | tb.Helper() 312 | if _, ok := s.LoadCounter(name); ok { 313 | tb.Errorf("gostats/mock: Counter (%q): expected Counter to not exist", name) 314 | } 315 | } 316 | 317 | // AssertGaugeNotExists asserts that Gauge name does not exist. 318 | func (s *Sink) AssertGaugeNotExists(tb testing.TB, name string) { 319 | tb.Helper() 320 | if _, ok := s.LoadGauge(name); ok { 321 | tb.Errorf("gostats/mock: Gauge (%q): expected Gauge to not exist", name) 322 | } 323 | } 324 | 325 | // AssertTimerNotExists asserts that Timer name does not exist. 326 | func (s *Sink) AssertTimerNotExists(tb testing.TB, name string) { 327 | tb.Helper() 328 | if _, ok := s.LoadTimer(name); ok { 329 | tb.Errorf("gostats/mock: Timer (%q): expected Timer to not exist", name) 330 | } 331 | } 332 | 333 | // AssertCounterCallCount asserts that Counter name was called exp times. 334 | func (s *Sink) AssertCounterCallCount(tb testing.TB, name string, exp int) { 335 | tb.Helper() 336 | v, ok := s.counters().Load(name) 337 | if !ok { 338 | tb.Errorf("gostats/mock: Counter (%q): not found in: %q", name, s.ListCounters()) 339 | return 340 | } 341 | p := v.(*entry) 342 | n := atomic.LoadInt64(&p.count) 343 | if n != int64(exp) { 344 | tb.Errorf("gostats/mock: Counter (%q) Call Count: Expected: %d Got: %d", 345 | name, exp, n) 346 | } 347 | } 348 | 349 | // AssertGaugeCallCount asserts that Gauge name was called exp times. 350 | func (s *Sink) AssertGaugeCallCount(tb testing.TB, name string, exp int) { 351 | tb.Helper() 352 | v, ok := s.gauges().Load(name) 353 | if !ok { 354 | tb.Errorf("gostats/mock: Gauge (%q): not found in: %q", name, s.ListGauges()) 355 | return 356 | } 357 | p := v.(*entry) 358 | n := atomic.LoadInt64(&p.count) 359 | if n != int64(exp) { 360 | tb.Errorf("gostats/mock: Gauge (%q) Call Count: Expected: %d Got: %d", 361 | name, exp, n) 362 | } 363 | } 364 | 365 | // AssertTimerCallCount asserts that Timer name was called exp times. 366 | func (s *Sink) AssertTimerCallCount(tb testing.TB, name string, exp int) { 367 | tb.Helper() 368 | v, ok := s.timers().Load(name) 369 | if !ok { 370 | tb.Errorf("gostats/mock: Timer (%q): not found in: %q", name, s.ListTimers()) 371 | return 372 | } 373 | p := v.(*entry) 374 | n := atomic.LoadInt64(&p.count) 375 | if n != int64(exp) { 376 | tb.Errorf("gostats/mock: Timer (%q) Call Count: Expected: %d Got: %d", 377 | name, exp, n) 378 | } 379 | } 380 | 381 | var ( 382 | _ testing.TB = (*fatalTest)(nil) 383 | _ testing.TB = (*fatalBench)(nil) 384 | ) 385 | 386 | type fatalTest testing.T 387 | 388 | func (t *fatalTest) Errorf(format string, args ...interface{}) { 389 | t.Fatalf(format, args...) 390 | } 391 | 392 | type fatalBench testing.B 393 | 394 | func (t *fatalBench) Errorf(format string, args ...interface{}) { 395 | t.Fatalf(format, args...) 396 | } 397 | 398 | // Fatal is a wrapper around *testing.T and *testing.B that causes Sink Assert* 399 | // methods to immediately fail a test and stop execution. Otherwise, the Assert 400 | // methods call tb.Errorf(), which marks the test as failed, but allows 401 | // execution to continue. 402 | // 403 | // Examples of Fatal() can be found in the sink test code. 404 | // 405 | // var sink Sink 406 | // var t *testing.T 407 | // sink.AssertCounterEquals(Must(t), "name", 1) 408 | func Fatal(tb testing.TB) testing.TB { 409 | switch t := tb.(type) { 410 | case *testing.T: 411 | return (*fatalTest)(t) 412 | case *testing.B: 413 | return (*fatalBench)(t) 414 | default: 415 | panic(fmt.Sprintf("invalid type for testing.TB: %T", tb)) 416 | } 417 | } 418 | 419 | // ParseTags extracts the name and tags from a statsd stat. 420 | // 421 | // Example of parsing tags and the stat name from a statsd stat: 422 | // 423 | // expected := map[string]string{ 424 | // "_f": "i", 425 | // "tag1": "value1", 426 | // } 427 | // name, tags := mock.ParseTags("prefix.c.___f=i.__tag1=value1") 428 | // if name != "panic.c" { 429 | // panic(fmt.Sprintf("Name: got: %q want: %q", name, "panic.c")) 430 | // } 431 | // if !reflect.DeepEqual(tags, expected) { 432 | // panic(fmt.Sprintf("Tags: got: %q want: %q", tags, expected)) 433 | // } 434 | func ParseTags(stat string) (string, map[string]string) { 435 | return tags.ParseTags(stat) 436 | } 437 | 438 | // SerializeTags serializes name and tags into a statsd stat. 439 | // 440 | // tags := map[string]string{ 441 | // "key_1": "val_1" 442 | // "key_2": "val_2" 443 | // } 444 | // s.AssertCounterExists(tb, SerializeTags("name", tags)) 445 | func SerializeTags(name string, tagsm map[string]string) string { 446 | return tags.SerializeTags(name, tagsm) 447 | } 448 | -------------------------------------------------------------------------------- /stats_test.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "context" 5 | crand "crypto/rand" 6 | "encoding/hex" 7 | "fmt" 8 | "math/rand" 9 | "reflect" 10 | "strconv" 11 | "strings" 12 | "sync" 13 | "testing" 14 | "time" 15 | 16 | tagspkg "github.com/lyft/gostats/internal/tags" 17 | "github.com/lyft/gostats/mock" 18 | ) 19 | 20 | // Ensure flushing and adding generators does not race 21 | func TestStats(_ *testing.T) { 22 | sink := &testStatSink{} 23 | store := NewStore(sink, true) 24 | 25 | scope := store.Scope("runtime") 26 | g := NewRuntimeStats(scope) 27 | var wg sync.WaitGroup 28 | wg.Add(2) 29 | 30 | go func() { 31 | store.AddStatGenerator(g) 32 | store.NewCounter("test") 33 | store.Flush() 34 | wg.Done() 35 | }() 36 | 37 | go func() { 38 | store.AddStatGenerator(g) 39 | store.NewCounter("test") 40 | store.Flush() 41 | wg.Done() 42 | }() 43 | 44 | wg.Wait() 45 | } 46 | 47 | // TestStatsStartContext ensures that a cancelled context cancels a 48 | // flushing goroutine. 49 | func TestStatsStartContext(_ *testing.T) { 50 | sink := &testStatSink{} 51 | store := NewStore(sink, true) 52 | 53 | ctx, cancel := context.WithCancel(context.Background()) 54 | tick := time.NewTicker(1 * time.Minute) 55 | 56 | wg := &sync.WaitGroup{} 57 | 58 | wg.Add(1) 59 | go func() { 60 | defer wg.Done() 61 | store.StartContext(ctx, tick) 62 | }() 63 | 64 | // now we cancel, and its ok to do this at any point - the 65 | // goroutine above could have started or not started, either case 66 | // is ok. 67 | cancel() 68 | 69 | wg.Wait() 70 | } 71 | 72 | // Ensure we create a counter and increment it for reserved tags 73 | func TestValidateTags(t *testing.T) { 74 | // Ensure we don't create a counter without reserved tags 75 | sink := &testStatSink{} 76 | store := NewStore(sink, true) 77 | store.NewCounter("test").Inc() 78 | store.Flush() 79 | 80 | expected := "test:1|c" 81 | output := sink.record 82 | if !strings.Contains(output, expected) && !strings.Contains(output, "reserved_tag") { 83 | t.Errorf("Expected without reserved tags: '%s' Got: '%s'", expected, output) 84 | } 85 | 86 | // A reserved tag should trigger adding the reserved_tag counter 87 | sink = &testStatSink{} 88 | store = NewStore(sink, true) 89 | store.NewCounterWithTags("test", map[string]string{"host": "i"}).Inc() 90 | store.Flush() 91 | 92 | expected = "test.__host=i:1|c" 93 | expectedReservedTag := "reserved_tag:1|c" 94 | output = sink.record 95 | if !strings.Contains(output, expected) && !strings.Contains(output, expectedReservedTag) { 96 | t.Errorf("Expected: '%s' and '%s', In: '%s'", expected, expectedReservedTag, output) 97 | } 98 | } 99 | 100 | // Ensure timers and timespans are working 101 | func TestTimer(t *testing.T) { 102 | testDuration := time.Duration(9800000) 103 | sink := &testStatSink{} 104 | store := NewStore(sink, true) 105 | store.NewTimer("test").AllocateSpan().CompleteWithDuration(testDuration) 106 | store.Flush() 107 | 108 | expected := "test:9800.000000|ms" 109 | timer := sink.record 110 | if !strings.Contains(timer, expected) { 111 | t.Error("wanted timer value of test:9800.000000|ms, got", timer) 112 | } 113 | } 114 | 115 | // Ensure millitimers and timespans are working 116 | func TestMilliTimer(t *testing.T) { 117 | testDuration := 420 * time.Millisecond 118 | sink := &testStatSink{} 119 | store := NewStore(sink, true) 120 | store.NewMilliTimer("test").AllocateSpan().CompleteWithDuration(testDuration) 121 | store.Flush() 122 | 123 | expected := "test:420.000000|ms" 124 | timer := sink.record 125 | if !strings.Contains(timer, expected) { 126 | t.Error("wanted timer value of test:420.000000|ms, got", timer) 127 | } 128 | } 129 | 130 | // Ensure 0 counters are not flushed 131 | func TestZeroCounters(t *testing.T) { 132 | sink := &testStatSink{} 133 | store := NewStore(sink, true) 134 | store.NewCounter("test") 135 | store.Flush() 136 | 137 | expected := "" 138 | counter := sink.record 139 | if counter != expected { 140 | t.Errorf("wanted %q got %q", expected, counter) 141 | } 142 | } 143 | 144 | func randomString(tb testing.TB, size int) string { 145 | b := make([]byte, hex.DecodedLen(size)) 146 | if _, err := crand.Read(b); err != nil { 147 | tb.Fatal(err) 148 | } 149 | return hex.EncodeToString(b) 150 | } 151 | 152 | func randomTagSet(t testing.TB, valPrefix string, size int) tagspkg.TagSet { 153 | s := make(tagspkg.TagSet, size) 154 | for i := 0; i < len(s); i++ { 155 | s[i] = tagspkg.NewTag(randomString(t, 32), fmt.Sprintf("%s%d", valPrefix, i)) 156 | } 157 | s.Sort() 158 | return s 159 | } 160 | 161 | func TestNewSubScope(t *testing.T) { 162 | s := randomTagSet(t, "x_", 20) 163 | for i := range s { 164 | s[i].Value += "|" // add an invalid char 165 | } 166 | m := make(map[string]string) 167 | for _, p := range s { 168 | m[p.Key] = p.Value 169 | } 170 | scope := newSubScope(nil, "name", m) 171 | 172 | expected := make(tagspkg.TagSet, len(s)) 173 | for i, p := range s { 174 | expected[i] = tagspkg.NewTag(p.Key, p.Value) 175 | } 176 | 177 | if !reflect.DeepEqual(scope.tags, expected) { 178 | t.Errorf("tags are not sorted by key: %+v", s) 179 | } 180 | for i, p := range expected { 181 | s := tagspkg.ReplaceChars(p.Value) 182 | if p.Value != s { 183 | t.Errorf("failed to replace invalid chars: %d: %+v", i, p) 184 | } 185 | } 186 | if scope.name != "name" { 187 | t.Errorf("wrong scope name: %s", scope.name) 188 | } 189 | } 190 | 191 | // Test that we never modify the tags map that is passed in 192 | func TestTagMapNotModified(t *testing.T) { 193 | type TagMethod func(scope Scope, name string, tags map[string]string) 194 | 195 | copyTags := func(tags map[string]string) map[string]string { 196 | orig := make(map[string]string, len(tags)) 197 | for k, v := range tags { 198 | orig[k] = v 199 | } 200 | return orig 201 | } 202 | 203 | scopeGenerators := map[string]func() Scope{ 204 | "statStore": func() Scope { return &statStore{} }, 205 | "subScope": func() Scope { return newSubScope(&statStore{}, "name", nil) }, 206 | } 207 | 208 | methodTestCases := map[string]TagMethod{ 209 | "ScopeWithTags": func(scope Scope, name string, tags map[string]string) { 210 | scope.ScopeWithTags(name, tags) 211 | }, 212 | "NewCounterWithTags": func(scope Scope, name string, tags map[string]string) { 213 | scope.NewCounterWithTags(name, tags) 214 | }, 215 | "NewPerInstanceCounter": func(scope Scope, name string, tags map[string]string) { 216 | scope.NewPerInstanceCounter(name, tags) 217 | }, 218 | "NewGaugeWithTags": func(scope Scope, name string, tags map[string]string) { 219 | scope.NewGaugeWithTags(name, tags) 220 | }, 221 | "NewPerInstanceGauge": func(scope Scope, name string, tags map[string]string) { 222 | scope.NewPerInstanceGauge(name, tags) 223 | }, 224 | "NewTimerWithTags": func(scope Scope, name string, tags map[string]string) { 225 | scope.NewTimerWithTags(name, tags) 226 | }, 227 | "NewPerInstanceTimer": func(scope Scope, name string, tags map[string]string) { 228 | scope.NewPerInstanceTimer(name, tags) 229 | }, 230 | } 231 | 232 | tagsTestCases := []map[string]string{ 233 | {}, // empty 234 | { 235 | "": "invalid_key", 236 | }, 237 | { 238 | "invalid_value": "", 239 | }, 240 | { 241 | "": "invalid_key", 242 | "invalid_value": "", 243 | }, 244 | { 245 | "_f": "i", 246 | }, 247 | { 248 | "": "invalid_key", 249 | "invalid_value": "", 250 | "_f": "i", 251 | }, 252 | { 253 | "": "invalid_key", 254 | "invalid_value": "", 255 | "_f": "value", 256 | "1": "1", 257 | }, 258 | { 259 | "": "invalid_key", 260 | "invalid_value": "", 261 | "1": "1", 262 | "2": "2", 263 | "3": "3", 264 | }, 265 | } 266 | 267 | for scopeName, newScope := range scopeGenerators { 268 | for methodName, method := range methodTestCases { 269 | t.Run(scopeName+"."+methodName, func(t *testing.T) { 270 | for _, orig := range tagsTestCases { 271 | tags := copyTags(orig) 272 | method(newScope(), "test", tags) 273 | if !reflect.DeepEqual(tags, orig) { 274 | t.Errorf("modified input map: %+v want: %+v", tags, orig) 275 | } 276 | } 277 | }) 278 | } 279 | } 280 | } 281 | 282 | func TestPerInstanceStats(t *testing.T) { 283 | testCases := []struct { 284 | expected string 285 | tags map[string]string 286 | }{ 287 | { 288 | expected: "name.___f=i", 289 | tags: map[string]string{}, // empty 290 | }, 291 | { 292 | expected: "name.___f=i", 293 | tags: map[string]string{ 294 | "": "invalid_key", 295 | }, 296 | }, 297 | { 298 | expected: "name.___f=i", 299 | tags: map[string]string{ 300 | "invalid_value": "", 301 | }, 302 | }, 303 | { 304 | expected: "name.___f=i", 305 | tags: map[string]string{ 306 | "_f": "i", 307 | }, 308 | }, 309 | { 310 | expected: "name.___f=xxx", 311 | tags: map[string]string{ 312 | "_f": "xxx", 313 | }, 314 | }, 315 | { 316 | expected: "name.___f=xxx", 317 | tags: map[string]string{ 318 | "": "invalid_key", 319 | "_f": "xxx", 320 | }, 321 | }, 322 | { 323 | expected: "name.___f=xxx", 324 | tags: map[string]string{ 325 | "invalid_value": "", 326 | "_f": "xxx", 327 | }, 328 | }, 329 | { 330 | expected: "name.___f=xxx", 331 | tags: map[string]string{ 332 | "invalid_value": "", 333 | "": "invalid_key", 334 | "_f": "xxx", 335 | }, 336 | }, 337 | { 338 | expected: "name.__1=1.___f=xxx", 339 | tags: map[string]string{ 340 | "invalid_value": "", 341 | "": "invalid_key", 342 | "_f": "xxx", 343 | "1": "1", 344 | }, 345 | }, 346 | { 347 | expected: "name.__1=1.___f=i", 348 | tags: map[string]string{ 349 | "1": "1", 350 | }, 351 | }, 352 | { 353 | expected: "name.__1=1.__2=2.___f=i", 354 | tags: map[string]string{ 355 | "1": "1", 356 | "2": "2", 357 | }, 358 | }, 359 | } 360 | 361 | testPerInstanceMethods := func(t *testing.T, setupScope func(Scope) Scope) { 362 | for _, x := range testCases { 363 | sink := mock.NewSink() 364 | scope := setupScope(&statStore{sink: sink}) 365 | 366 | scope.NewPerInstanceCounter("name", x.tags).Inc() 367 | scope.NewPerInstanceGauge("name", x.tags).Inc() 368 | scope.NewPerInstanceTimer("name", x.tags).AddValue(1) 369 | scope.Store().Flush() 370 | 371 | for key := range sink.Counters() { 372 | if key != x.expected { 373 | t.Errorf("Counter (%+v): got: %q want: %q", x, key, x.expected) 374 | } 375 | break 376 | } 377 | 378 | for key := range sink.Gauges() { 379 | if key != x.expected { 380 | t.Errorf("Gauge (%+v): got: %q want: %q", x, key, x.expected) 381 | } 382 | break 383 | } 384 | 385 | for key := range sink.Timers() { 386 | if key != x.expected { 387 | t.Errorf("Timer (%+v): got: %q want: %q", x, key, x.expected) 388 | } 389 | break 390 | } 391 | } 392 | } 393 | 394 | t.Run("StatsStore", func(t *testing.T) { 395 | testPerInstanceMethods(t, func(scope Scope) Scope { return scope }) 396 | }) 397 | 398 | t.Run("SubScope", func(t *testing.T) { 399 | // Add sub-scope prefix to the name 400 | for i, x := range testCases { 401 | testCases[i].expected = "x." + x.expected 402 | } 403 | 404 | testPerInstanceMethods(t, func(scope Scope) Scope { 405 | return scope.Scope("x") 406 | }) 407 | }) 408 | } 409 | 410 | func BenchmarkStore_MutexContention(b *testing.B) { 411 | s := NewStore(nullSink{}, false) 412 | t := time.NewTicker(500 * time.Microsecond) // we want flush to contend with accessing metrics 413 | defer t.Stop() 414 | go s.Start(t) 415 | 416 | b.ResetTimer() 417 | for i := 0; i < b.N; i++ { 418 | bmID := strconv.Itoa(rand.Intn(1000)) 419 | c := s.NewCounter(bmID) 420 | c.Inc() 421 | _ = c.Value() 422 | } 423 | } 424 | 425 | func BenchmarkStore_NewCounterWithTags(b *testing.B) { 426 | s := NewStore(nullSink{}, false) 427 | t := time.NewTicker(time.Hour) // don't flush 428 | defer t.Stop() 429 | go s.Start(t) 430 | tags := map[string]string{ 431 | "tag1": "val1", 432 | "tag2": "val2", 433 | "tag3": "val3", 434 | "tag4": "val4", 435 | "tag5": "val5", 436 | } 437 | b.ResetTimer() 438 | for i := 0; i < b.N; i++ { 439 | s.NewCounterWithTags("counter_name", tags) 440 | } 441 | } 442 | 443 | func initBenchScope() (scope Scope, childTags map[string]string) { 444 | s := NewStore(nullSink{}, false) 445 | 446 | scopeTags := make(map[string]string, 5) 447 | childTags = make(map[string]string, 5) 448 | 449 | for i := 0; i < 5; i++ { 450 | tag := fmt.Sprintf("%dtag", i) 451 | val := fmt.Sprintf("%dval", i) 452 | scopeTags[tag] = val 453 | childTags["c"+tag] = "c" + val 454 | } 455 | 456 | scope = s.ScopeWithTags("scope", scopeTags) 457 | return 458 | } 459 | 460 | func BenchmarkStore_ScopeWithTags(b *testing.B) { 461 | scope, childTags := initBenchScope() 462 | b.ResetTimer() 463 | for i := 0; i < b.N; i++ { 464 | scope.NewCounterWithTags("counter_name", childTags) 465 | } 466 | } 467 | 468 | func BenchmarkStore_ScopeNoTags(b *testing.B) { 469 | scope, _ := initBenchScope() 470 | b.ResetTimer() 471 | for i := 0; i < b.N; i++ { 472 | scope.NewCounterWithTags("counter_name", nil) 473 | } 474 | } 475 | 476 | func BenchmarkParallelCounter(b *testing.B) { 477 | const N = 1000 478 | keys := make([]string, N) 479 | for i := 0; i < len(keys); i++ { 480 | keys[i] = randomString(b, 32) 481 | } 482 | 483 | s := NewStore(nullSink{}, false) 484 | t := time.NewTicker(time.Hour) // don't flush 485 | defer t.Stop() // never sends 486 | go s.Start(t) 487 | 488 | b.ResetTimer() 489 | b.RunParallel(func(pb *testing.PB) { 490 | n := 0 491 | for pb.Next() { 492 | s.NewCounter(keys[n%N]).Inc() 493 | } 494 | }) 495 | } 496 | 497 | func BenchmarkStoreNewPerInstanceCounter(b *testing.B) { 498 | b.Run("HasTag", func(b *testing.B) { 499 | var store statStore 500 | tags := map[string]string{ 501 | "1": "1", 502 | "2": "2", 503 | "3": "3", 504 | "_f": "xxx", 505 | } 506 | for i := 0; i < b.N; i++ { 507 | store.NewPerInstanceCounter("name", tags) 508 | } 509 | }) 510 | 511 | b.Run("MissingTag", func(b *testing.B) { 512 | var store statStore 513 | tags := map[string]string{ 514 | "1": "1", 515 | "2": "2", 516 | "3": "3", 517 | "4": "4", 518 | } 519 | for i := 0; i < b.N; i++ { 520 | store.NewPerInstanceCounter("name", tags) 521 | } 522 | }) 523 | } 524 | -------------------------------------------------------------------------------- /internal/tags/tags.go: -------------------------------------------------------------------------------- 1 | package tags 2 | 3 | import ( 4 | "sort" 5 | "strings" 6 | "unsafe" 7 | ) 8 | 9 | // A Tag is a Key/Value statsd tag. 10 | type Tag struct { 11 | Key string 12 | Value string 13 | } 14 | 15 | // NewTag returns a new Tag with any invalid chars in the value replaced. 16 | func NewTag(key, value string) Tag { 17 | return Tag{Key: key, Value: ReplaceChars(value)} 18 | } 19 | 20 | // A TagSet is a collection of Tags. All methods apart from Sort() require the 21 | // TagSet to be sorted. Tags with empty keys or values should never be inserted 22 | // into the TagSet since most methods rely on all the tags being valid. 23 | type TagSet []Tag 24 | 25 | // NewTagSet returns a new TagSet from the tags map. 26 | func NewTagSet(tags map[string]string) TagSet { 27 | a := make(TagSet, 0, len(tags)) 28 | for k, v := range tags { 29 | if k != "" && v != "" { 30 | a = append(a, NewTag(k, v)) 31 | } 32 | } 33 | a.Sort() 34 | return a 35 | } 36 | 37 | func (t TagSet) Len() int { return len(t) } 38 | func (t TagSet) Swap(i, j int) { t[i], t[j] = t[j], t[i] } 39 | func (t TagSet) Less(i, j int) bool { return t[i].Key < t[j].Key } 40 | 41 | // cas performs a compare and swap and is inlined into Sort() 42 | // check with: `go build -gcflags='-m'`. 43 | func (t TagSet) cas(i, j int) { 44 | if t[i].Key > t[j].Key { 45 | t.Swap(i, j) 46 | } 47 | } 48 | 49 | // Sort sorts the TagSet in place and is optimized for small (N <= 8) TagSets. 50 | func (t TagSet) Sort() { 51 | // network sort generated with: https://pages.ripco.net/~jgamble/nw.html 52 | // using "best": https://metacpan.org/pod/Algorithm::Networksort::Best 53 | // 54 | // example query (N=6): 55 | // http://jgamble.ripco.net/cgi-bin/nw.cgi?inputs=6&algorithm=best&output=macro 56 | // 57 | // all cas() methods are inlined, check with: `go build -gcflags='-m'` 58 | switch len(t) { 59 | case 0, 1: 60 | return 61 | case 2: 62 | t.cas(0, 1) 63 | case 3: 64 | t.cas(1, 2) 65 | t.cas(0, 2) 66 | t.cas(0, 1) 67 | case 4: 68 | t.cas(0, 1) 69 | t.cas(2, 3) 70 | t.cas(0, 2) 71 | t.cas(1, 3) 72 | t.cas(1, 2) 73 | case 5: 74 | t.cas(0, 1) 75 | t.cas(3, 4) 76 | t.cas(2, 4) 77 | t.cas(2, 3) 78 | t.cas(0, 3) 79 | t.cas(0, 2) 80 | t.cas(1, 4) 81 | t.cas(1, 3) 82 | t.cas(1, 2) 83 | case 6: 84 | t.cas(1, 2) 85 | t.cas(0, 2) 86 | t.cas(0, 1) 87 | t.cas(4, 5) 88 | t.cas(3, 5) 89 | t.cas(3, 4) 90 | t.cas(0, 3) 91 | t.cas(1, 4) 92 | t.cas(2, 5) 93 | t.cas(2, 4) 94 | t.cas(1, 3) 95 | t.cas(2, 3) 96 | case 7: 97 | t.cas(1, 2) 98 | t.cas(0, 2) 99 | t.cas(0, 1) 100 | t.cas(3, 4) 101 | t.cas(5, 6) 102 | t.cas(3, 5) 103 | t.cas(4, 6) 104 | t.cas(4, 5) 105 | t.cas(0, 4) 106 | t.cas(0, 3) 107 | t.cas(1, 5) 108 | t.cas(2, 6) 109 | t.cas(2, 5) 110 | t.cas(1, 3) 111 | t.cas(2, 4) 112 | t.cas(2, 3) 113 | case 8: 114 | t.cas(0, 1) 115 | t.cas(2, 3) 116 | t.cas(0, 2) 117 | t.cas(1, 3) 118 | t.cas(1, 2) 119 | t.cas(4, 5) 120 | t.cas(6, 7) 121 | t.cas(4, 6) 122 | t.cas(5, 7) 123 | t.cas(5, 6) 124 | t.cas(0, 4) 125 | t.cas(1, 5) 126 | t.cas(1, 4) 127 | t.cas(2, 6) 128 | t.cas(3, 7) 129 | t.cas(3, 6) 130 | t.cas(2, 4) 131 | t.cas(3, 5) 132 | t.cas(3, 4) 133 | default: 134 | sort.Sort(t) 135 | } 136 | } 137 | 138 | // Search is the same as sort.Search() but is optimized for our use case. 139 | func (t TagSet) Search(key string) int { 140 | i, j := 0, len(t) 141 | for i < j { 142 | h := (i + j) / 2 143 | if t[h].Key < key { 144 | i = h + 1 145 | } else { 146 | j = h 147 | } 148 | } 149 | return i 150 | } 151 | 152 | // Contains returns if the TagSet contains key. 153 | func (t TagSet) Contains(key string) bool { 154 | if len(t) == 0 { 155 | return false 156 | } 157 | i := t.Search(key) 158 | return i < len(t) && t[i].Key == key 159 | } 160 | 161 | // Insert inserts Tag p into TagSet t, returning a copy unless Tag p already 162 | // existed. 163 | func (t TagSet) Insert(p Tag) TagSet { 164 | if len(t) == 0 { 165 | return TagSet{p} 166 | } 167 | 168 | i := t.Search(p.Key) 169 | if i < len(t) && t[i].Key == p.Key { 170 | if t[i].Value == p.Value { 171 | return t // no change 172 | } 173 | a := make(TagSet, len(t)) 174 | copy(a, t) 175 | a[i].Value = p.Value 176 | return a // exists 177 | } 178 | 179 | // we're modifying the set - make a copy 180 | a := make(TagSet, len(t)+1) 181 | copy(a[:i], t[:i]) 182 | a[i] = p 183 | copy(a[i+1:], t[i:]) 184 | return a 185 | } 186 | 187 | // MergeTags returns a TagSet that is the union of subScope's tags and the 188 | // provided tags map. If any keys overlap the values from the provided map 189 | // are used. 190 | func (t TagSet) MergeTags(tags map[string]string) TagSet { 191 | switch len(tags) { 192 | case 0: 193 | return t 194 | case 1: 195 | // optimize for the common case of there only being one tag 196 | return mergeOneTag(t, tags) 197 | default: 198 | // write tags to the end of the scratch slice 199 | scratch := make(TagSet, len(t)+len(tags)) 200 | a := scratch[len(t):] 201 | i := 0 202 | for k, v := range tags { 203 | if k != "" && v != "" { 204 | a[i] = NewTag(k, v) 205 | i++ 206 | } 207 | } 208 | a = a[:i] 209 | a.Sort() 210 | 211 | if len(t) == 0 { 212 | return a 213 | } 214 | return mergeTagSets(t, a, scratch) 215 | } 216 | } 217 | 218 | // mergeOneTag is an optimized for inserting 1 tag and will panic otherwise. 219 | func mergeOneTag(set TagSet, tags map[string]string) TagSet { 220 | if len(tags) != 1 { 221 | panic("invalid usage") 222 | } 223 | var p Tag 224 | for k, v := range tags { 225 | p = NewTag(k, v) 226 | break 227 | } 228 | if p.Key == "" || p.Value == "" { 229 | return set 230 | } 231 | return set.Insert(p) 232 | } 233 | 234 | // MergePerInstanceTags returns a TagSet that is the union of subScope's 235 | // tags and the provided tags map with. If any keys overlap the values from 236 | // the provided map are used. 237 | // 238 | // The returned TagSet will have a per-instance key ("_f") and if neither the 239 | // subScope or tags have this key it's value will be the default per-instance 240 | // value ("i"). 241 | // 242 | // The method does not optimize for the case where there is only one tag 243 | // because it is used less frequently. 244 | func (t TagSet) MergePerInstanceTags(tags map[string]string) TagSet { 245 | if len(tags) == 0 { 246 | if t.Contains("_f") { 247 | return t 248 | } 249 | // create copy with the per-instance tag 250 | return t.Insert(Tag{Key: "_f", Value: "i"}) 251 | } 252 | 253 | // write tags to the end of scratch slice 254 | scratch := make(TagSet, len(t)+len(tags)+1) 255 | a := scratch[len(t):] 256 | i := 0 257 | // add the default per-instance tag if not present 258 | if tags["_f"] == "" && !t.Contains("_f") { 259 | a[i] = Tag{Key: "_f", Value: "i"} 260 | i++ 261 | } 262 | for k, v := range tags { 263 | if k != "" && v != "" { 264 | a[i] = NewTag(k, v) 265 | i++ 266 | } 267 | } 268 | a = a[:i] 269 | a.Sort() 270 | 271 | if len(t) == 0 { 272 | return a 273 | } 274 | return mergeTagSets(t, a, scratch) 275 | } 276 | 277 | // mergeTagSets merges s1 into s2 and stores the result in scratch. Both s1 and 278 | // s2 must be sorted and s2 can be a sub-slice of scratch if it is located at 279 | // the tail of the slice (this allows us to allocate only one slice). 280 | func mergeTagSets(s1, s2, scratch TagSet) TagSet { 281 | a := scratch 282 | i, j, k := 0, 0, 0 283 | for ; i < len(s1) && j < len(s2) && k < len(a); k++ { 284 | if s1[i].Key == s2[j].Key { 285 | a[k] = s2[j] 286 | i++ 287 | j++ 288 | } else if s1[i].Key < s2[j].Key { 289 | a[k] = s1[i] 290 | i++ 291 | } else { 292 | a[k] = s2[j] 293 | j++ 294 | } 295 | } 296 | if i < len(s1) { 297 | k += copy(a[k:], s1[i:]) 298 | } 299 | if j < len(s2) { 300 | k += copy(a[k:], s2[j:]) 301 | } 302 | return a[:k] 303 | } 304 | 305 | // Serialize serializes name and tags into a statsd stat. Note: the TagSet 306 | // t must be sorted and have clean tag keys and values. 307 | func (t TagSet) Serialize(name string) string { 308 | // TODO: panic if the set isn't sorted? 309 | 310 | const prefix = ".__" 311 | const sep = "=" 312 | 313 | if len(t) == 0 { 314 | return name 315 | } 316 | 317 | n := (len(prefix)+len(sep))*len(t) + len(name) 318 | for _, p := range t { 319 | n += len(p.Key) + len(p.Value) 320 | } 321 | 322 | // CEV: this is same as strings.Builder, but is faster and simpler. 323 | b := make([]byte, 0, n) 324 | b = append(b, name...) 325 | for _, p := range t { 326 | b = append(b, prefix...) 327 | b = append(b, p.Key...) 328 | b = append(b, sep...) 329 | b = append(b, p.Value...) 330 | } 331 | return *(*string)(unsafe.Pointer(&b)) 332 | } 333 | 334 | // SerializeTags serializes name and tags into a statsd stat. 335 | func SerializeTags(name string, tags map[string]string) string { 336 | const prefix = ".__" 337 | const sep = "=" 338 | 339 | // discard pairs where the tag or value is an empty string 340 | numValid := len(tags) 341 | for k, v := range tags { 342 | if k == "" || v == "" { 343 | numValid-- 344 | } 345 | } 346 | 347 | switch numValid { 348 | case 0: 349 | return name 350 | case 1: 351 | var t0 Tag 352 | for k, v := range tags { 353 | if k != "" && v != "" { 354 | t0 = NewTag(k, v) 355 | break 356 | } 357 | } 358 | return name + prefix + t0.Key + sep + t0.Value 359 | case 2: 360 | var t0, t1 Tag 361 | for k, v := range tags { 362 | if k == "" || v == "" { 363 | continue 364 | } 365 | t1 = t0 366 | t0 = NewTag(k, v) 367 | } 368 | if t0.Key > t1.Key { 369 | t0, t1 = t1, t0 370 | } 371 | return name + prefix + t0.Key + sep + t0.Value + 372 | prefix + t1.Key + sep + t1.Value 373 | case 3: 374 | var t0, t1, t2 Tag 375 | for k, v := range tags { 376 | if k == "" || v == "" { 377 | continue 378 | } 379 | t2 = t1 380 | t1 = t0 381 | t0 = NewTag(k, v) 382 | } 383 | if t1.Key > t2.Key { 384 | t1, t2 = t2, t1 385 | } 386 | if t0.Key > t2.Key { 387 | t0, t2 = t2, t0 388 | } 389 | if t0.Key > t1.Key { 390 | t0, t1 = t1, t0 391 | } 392 | return name + prefix + t0.Key + sep + t0.Value + 393 | prefix + t1.Key + sep + t1.Value + 394 | prefix + t2.Key + sep + t2.Value 395 | case 4: 396 | var t0, t1, t2, t3 Tag 397 | for k, v := range tags { 398 | if k == "" || v == "" { 399 | continue 400 | } 401 | t3 = t2 402 | t2 = t1 403 | t1 = t0 404 | t0 = NewTag(k, v) 405 | } 406 | if t0.Key > t1.Key { 407 | t0, t1 = t1, t0 408 | } 409 | if t2.Key > t3.Key { 410 | t2, t3 = t3, t2 411 | } 412 | if t0.Key > t2.Key { 413 | t0, t2 = t2, t0 414 | } 415 | if t1.Key > t3.Key { 416 | t1, t3 = t3, t1 417 | } 418 | if t1.Key > t2.Key { 419 | t1, t2 = t2, t1 420 | } 421 | return name + prefix + t0.Key + sep + t0.Value + 422 | prefix + t1.Key + sep + t1.Value + 423 | prefix + t2.Key + sep + t2.Value + 424 | prefix + t3.Key + sep + t3.Value 425 | default: 426 | // n stores the length of the serialized name + tags 427 | n := (len(prefix) + len(sep)) * numValid 428 | n += len(name) 429 | 430 | pairs := make(TagSet, 0, numValid) 431 | for k, v := range tags { 432 | if k == "" || v == "" { 433 | continue 434 | } 435 | n += len(k) + len(v) 436 | pairs = append(pairs, NewTag(k, v)) 437 | } 438 | sort.Sort(pairs) 439 | 440 | // CEV: this is same as strings.Builder, but works with go1.9 and earlier 441 | b := make([]byte, 0, n) 442 | b = append(b, name...) 443 | for _, tag := range pairs { 444 | b = append(b, prefix...) 445 | b = append(b, tag.Key...) 446 | b = append(b, sep...) 447 | b = append(b, tag.Value...) 448 | } 449 | return *(*string)(unsafe.Pointer(&b)) 450 | } 451 | } 452 | 453 | // ReplaceChars replaces any invalid chars ([.:|]) in value s with '_'. 454 | func ReplaceChars(s string) string { 455 | var buf []byte // lazily allocated 456 | for i := 0; i < len(s); i++ { 457 | switch s[i] { 458 | case '.', ':', '|': 459 | if buf == nil { 460 | buf = []byte(s) 461 | } 462 | buf[i] = '_' 463 | } 464 | } 465 | if buf == nil { 466 | return s 467 | } 468 | return *(*string)(unsafe.Pointer(&buf)) 469 | } 470 | 471 | // removeStatValue removes the value from a stat line 472 | func removeStatValue(s string) string { 473 | i := strings.IndexByte(s, ':') 474 | if i == -1 { 475 | return s 476 | } 477 | n := i 478 | i++ 479 | for i < len(s) { 480 | if s[i] == ':' { 481 | n = i 482 | } 483 | i++ 484 | } 485 | return s[:n] 486 | } 487 | 488 | // ParseTags parses the statsd stat name and tags (if any) from stat. 489 | func ParseTags(stat string) (string, map[string]string) { 490 | const sep = ".__" 491 | 492 | // Remove the value, if any. This allows passing full stat 493 | // lines in wire form as they would be emitted by the sink. 494 | stat = removeStatValue(stat) 495 | 496 | o := strings.Index(stat, sep) 497 | if o == -1 { 498 | return stat, nil // no tags 499 | } 500 | name := stat[:o] 501 | 502 | // consume first sep 503 | s := stat[o+len(sep):] 504 | 505 | n := strings.Count(s, sep) 506 | tags := make(map[string]string, n+1) 507 | 508 | for ; n > 0; n-- { 509 | m := strings.Index(s, sep) 510 | if m < 0 { 511 | break 512 | } 513 | a := s[:m] 514 | if o := strings.IndexByte(a, '='); o != -1 { 515 | tags[a[:o]] = a[o+1:] 516 | } 517 | s = s[m+len(sep):] 518 | } 519 | if o := strings.IndexByte(s, '='); o != -1 { 520 | tags[s[:o]] = s[o+1:] 521 | } 522 | return name, tags 523 | } 524 | 525 | // ParseTagSet parses the statsd stat name and tags (if any) from stat. It is 526 | // like ParseTags, but returns a TagSet instead of a map[string]string. 527 | func ParseTagSet(stat string) (string, TagSet) { 528 | const sep = ".__" 529 | 530 | // Remove the value, if any. This allows passing full stat 531 | // lines in wire form as they would be emitted by the sink. 532 | stat = removeStatValue(stat) 533 | 534 | o := strings.Index(stat, sep) 535 | if o == -1 { 536 | return stat, nil // no tags 537 | } 538 | name := stat[:o] 539 | 540 | // consume first sep 541 | s := stat[o+len(sep):] 542 | 543 | n := strings.Count(s, sep) 544 | tags := make(TagSet, n+1) 545 | 546 | i := 0 547 | for n > 0 { 548 | m := strings.Index(s, sep) 549 | if m < 0 { 550 | break 551 | } 552 | // TODO: handle malformed stats ??? 553 | a := s[:m] 554 | if o := strings.IndexByte(a, '='); o != -1 { 555 | tags[i] = Tag{ 556 | Key: a[:o], 557 | Value: a[o+1:], // don't clean the stat 558 | } 559 | i++ 560 | } 561 | s = s[m+len(sep):] 562 | n-- 563 | } 564 | if o := strings.IndexByte(s, '='); o != -1 { 565 | tags[i] = Tag{ 566 | Key: s[:o], 567 | Value: s[o+1:], 568 | } 569 | i++ 570 | } 571 | tags = tags[:i] 572 | tags.Sort() 573 | return name, tags 574 | } 575 | -------------------------------------------------------------------------------- /stats.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "context" 5 | "strconv" 6 | "sync" 7 | "sync/atomic" 8 | "time" 9 | 10 | tagspkg "github.com/lyft/gostats/internal/tags" 11 | ) 12 | 13 | // A Store holds statistics. 14 | // There are two options when creating a new store: 15 | // 16 | // create a store backed by a tcp_sink to statsd 17 | // s := stats.NewDefaultStore() 18 | // create a store with a user provided Sink 19 | // s := stats.NewStore(sink, true) 20 | // 21 | // Currently that only backing store supported is statsd via a TCP sink, https://github.com/lyft/gostats/blob/master/tcp_sink.go. 22 | // However, implementing other Sinks (https://github.com/lyft/gostats/blob/master/sink.go) should be simple. 23 | // 24 | // A store holds Counters, Gauges, and Timers. You can add unscoped Counters, Gauges, and Timers to the store 25 | // with: 26 | // 27 | // s := stats.NewDefaultStore() 28 | // c := s.New[Counter|Gauge|Timer]("name") 29 | type Store interface { 30 | // Flush Counters and Gauges to the Sink attached to the Store. 31 | // To flush the store at a regular interval call the 32 | // Start(*time.Ticker) 33 | // method on it. 34 | // 35 | // The store will flush either at the regular interval, or whenever 36 | // Flush() 37 | // is called. Whenever the store is flushed, 38 | // the store will call 39 | // GenerateStats() 40 | // on all of its stat generators, 41 | // and flush all the Counters and Gauges registered with it. 42 | Flush() 43 | 44 | // Start a timer for periodic stat flushes. This is a blocking 45 | // call and should be called in a goroutine. 46 | Start(*time.Ticker) 47 | 48 | // StartContext starts a timer for periodic stat flushes. This is 49 | // a blocking call and should be called in a goroutine. 50 | // 51 | // If the passed-in context is cancelled, then this call 52 | // exits. Flush will be called on exit. 53 | StartContext(context.Context, *time.Ticker) 54 | 55 | // Add a StatGenerator to the Store that programatically generates stats. 56 | AddStatGenerator(StatGenerator) 57 | Scope 58 | } 59 | 60 | // A Scope namespaces Statistics. 61 | // 62 | // store := stats.NewDefaultStore() 63 | // scope := stats.Scope("service") 64 | // // the following counter will be emitted at the stats tree rooted at `service`. 65 | // c := scope.NewCounter("success") 66 | // 67 | // Additionally you can create subscopes: 68 | // 69 | // store := stats.NewDefaultStore() 70 | // scope := stats.Scope("service") 71 | // networkScope := scope.Scope("network") 72 | // // the following counter will be emitted at the stats tree rooted at service.network. 73 | // c := networkScope.NewCounter("requests") 74 | type Scope interface { 75 | // Scope creates a subscope. 76 | Scope(name string) Scope 77 | 78 | // ScopeWithTags creates a subscope with Tags to a store or scope. All child scopes and metrics 79 | // will inherit these tags by default. 80 | ScopeWithTags(name string, tags map[string]string) Scope 81 | 82 | // Store returns the Scope's backing Store. 83 | Store() Store 84 | 85 | // NewCounter adds a Counter to a store, or a scope. 86 | NewCounter(name string) Counter 87 | 88 | // NewCounterWithTags adds a Counter with Tags to a store, or a scope. 89 | NewCounterWithTags(name string, tags map[string]string) Counter 90 | 91 | // NewPerInstanceCounter adds a Per instance Counter with optional Tags to a store, or a scope. 92 | NewPerInstanceCounter(name string, tags map[string]string) Counter 93 | 94 | // NewGauge adds a Gauge to a store, or a scope. 95 | NewGauge(name string) Gauge 96 | 97 | // NewGaugeWithTags adds a Gauge with Tags to a store, or a scope. 98 | NewGaugeWithTags(name string, tags map[string]string) Gauge 99 | 100 | // NewPerInstanceGauge adds a Per instance Gauge with optional Tags to a store, or a scope. 101 | NewPerInstanceGauge(name string, tags map[string]string) Gauge 102 | 103 | // NewTimer adds a Timer to a store, or a scope that uses microseconds as its unit. 104 | NewTimer(name string) Timer 105 | 106 | // NewTimerWithTags adds a Timer with Tags to a store, or a scope with Tags that uses microseconds as its unit. 107 | NewTimerWithTags(name string, tags map[string]string) Timer 108 | 109 | // NewPerInstanceTimer adds a Per instance Timer with optional Tags to a store, or a scope that uses microseconds as its unit. 110 | NewPerInstanceTimer(name string, tags map[string]string) Timer 111 | 112 | // NewMilliTimer adds a Timer to a store, or a scope that uses milliseconds as its unit. 113 | NewMilliTimer(name string) Timer 114 | 115 | // NewMilliTimerWithTags adds a Timer with Tags to a store, or a scope with Tags that uses milliseconds as its unit. 116 | NewMilliTimerWithTags(name string, tags map[string]string) Timer 117 | 118 | // NewPerInstanceMilliTimer adds a Per instance Timer with optional Tags to a store, or a scope that uses milliseconds as its unit. 119 | NewPerInstanceMilliTimer(name string, tags map[string]string) Timer 120 | } 121 | 122 | // A Counter is an always incrementing stat. 123 | type Counter interface { 124 | // Add increments the Counter by the argument's value. 125 | Add(uint64) 126 | 127 | // Inc increments the Counter by 1. 128 | Inc() 129 | 130 | // Set sets an internal counter value which will be written in the next flush. 131 | // Its use is discouraged as it may break the counter's "always incrementing" semantics. 132 | Set(uint64) 133 | 134 | // String returns the current value of the Counter as a string. 135 | String() string 136 | 137 | // Value returns the current value of the Counter as a uint64. 138 | Value() uint64 139 | } 140 | 141 | // A Gauge is a stat that can increment and decrement. 142 | type Gauge interface { 143 | // Add increments the Gauge by the argument's value. 144 | Add(uint64) 145 | 146 | // Sub decrements the Gauge by the argument's value. 147 | Sub(uint64) 148 | 149 | // Inc increments the Gauge by 1. 150 | Inc() 151 | 152 | // Dec decrements the Gauge by 1. 153 | Dec() 154 | 155 | // Set sets the Gauge to a value. 156 | Set(uint64) 157 | 158 | // String returns the current value of the Gauge as a string. 159 | String() string 160 | 161 | // Value returns the current value of the Gauge as a uint64. 162 | Value() uint64 163 | } 164 | 165 | // A Timer is used to flush timing statistics. 166 | type Timer interface { 167 | // AddValue flushs the timer with the argument's value. 168 | AddValue(float64) 169 | 170 | // AddDuration emits the duration as a timing measurement. 171 | AddDuration(time.Duration) 172 | 173 | // AllocateSpan allocates a Timespan. 174 | AllocateSpan() Timespan 175 | } 176 | 177 | // A Timespan is used to measure spans of time. 178 | // They measure time from the time they are allocated by a Timer with 179 | // 180 | // AllocateSpan() 181 | // 182 | // until they call 183 | // 184 | // Complete() 185 | // 186 | // or 187 | // 188 | // CompleteWithDuration(time.Duration) 189 | // 190 | // When either function is called the timespan is flushed. 191 | // When Complete is called the timespan is flushed. 192 | // 193 | // A Timespan can be flushed at function 194 | // return by calling Complete with golang's defer statement. 195 | type Timespan interface { 196 | // End the Timespan and flush it. 197 | Complete() time.Duration 198 | 199 | // End the Timespan and flush it. Adds additional time.Duration to the measured time 200 | CompleteWithDuration(time.Duration) 201 | } 202 | 203 | // A StatGenerator can be used to programatically generate stats. 204 | // StatGenerators are added to a store via 205 | // 206 | // AddStatGenerator(StatGenerator) 207 | // 208 | // An example is https://github.com/lyft/gostats/blob/master/runtime.go. 209 | type StatGenerator interface { 210 | // Runs the StatGenerator to generate Stats. 211 | GenerateStats() 212 | } 213 | 214 | // NewStore returns an Empty store that flushes to Sink passed as an argument. 215 | // Note: the export argument is unused. 216 | func NewStore(sink Sink, _ bool) Store { 217 | return &statStore{sink: sink} 218 | } 219 | 220 | // NewDefaultStore returns a Store with a TCP statsd sink, and a running flush timer. 221 | func NewDefaultStore() Store { 222 | var newStore Store 223 | settings := GetSettings() 224 | if !settings.UseStatsd { 225 | if settings.LoggingSinkDisabled { 226 | newStore = NewStore(NewNullSink(), false) 227 | } else { 228 | newStore = NewStore(NewLoggingSink(), false) 229 | } 230 | go newStore.Start(time.NewTicker(10 * time.Second)) 231 | } else { 232 | newStore = NewStore(NewTCPStatsdSink(), false) 233 | go newStore.Start(time.NewTicker(time.Duration(settings.FlushIntervalS) * time.Second)) 234 | } 235 | return newStore 236 | } 237 | 238 | type counter struct { 239 | currentValue uint64 240 | lastSentValue uint64 241 | } 242 | 243 | func (c *counter) Add(delta uint64) { 244 | atomic.AddUint64(&c.currentValue, delta) 245 | } 246 | 247 | func (c *counter) Set(value uint64) { 248 | atomic.StoreUint64(&c.currentValue, value) 249 | } 250 | 251 | func (c *counter) Inc() { 252 | c.Add(1) 253 | } 254 | 255 | func (c *counter) Value() uint64 { 256 | return atomic.LoadUint64(&c.currentValue) 257 | } 258 | 259 | func (c *counter) String() string { 260 | return strconv.FormatUint(c.Value(), 10) 261 | } 262 | 263 | func (c *counter) latch() uint64 { 264 | value := c.Value() 265 | lastSent := atomic.SwapUint64(&c.lastSentValue, value) 266 | return value - lastSent 267 | } 268 | 269 | type gauge struct { 270 | value uint64 271 | } 272 | 273 | func (c *gauge) String() string { 274 | return strconv.FormatUint(c.Value(), 10) 275 | } 276 | 277 | func (c *gauge) Add(value uint64) { 278 | atomic.AddUint64(&c.value, value) 279 | } 280 | 281 | func (c *gauge) Sub(value uint64) { 282 | atomic.AddUint64(&c.value, ^(value - 1)) 283 | } 284 | 285 | func (c *gauge) Inc() { 286 | c.Add(1) 287 | } 288 | 289 | func (c *gauge) Dec() { 290 | c.Sub(1) 291 | } 292 | 293 | func (c *gauge) Set(value uint64) { 294 | atomic.StoreUint64(&c.value, value) 295 | } 296 | 297 | func (c *gauge) Value() uint64 { 298 | return atomic.LoadUint64(&c.value) 299 | } 300 | 301 | type timer struct { 302 | base time.Duration 303 | name string 304 | sink Sink 305 | } 306 | 307 | func (t *timer) time(dur time.Duration) { 308 | t.AddDuration(dur) 309 | } 310 | 311 | func (t *timer) AddDuration(dur time.Duration) { 312 | t.AddValue(float64(dur / t.base)) 313 | } 314 | 315 | func (t *timer) AddValue(value float64) { 316 | t.sink.FlushTimer(t.name, value) 317 | } 318 | 319 | func (t *timer) AllocateSpan() Timespan { 320 | return ×pan{timer: t, start: time.Now()} 321 | } 322 | 323 | type timespan struct { 324 | timer *timer 325 | start time.Time 326 | } 327 | 328 | func (ts *timespan) Complete() time.Duration { 329 | d := time.Since(ts.start) 330 | ts.timer.time(d) 331 | return d 332 | } 333 | 334 | func (ts *timespan) CompleteWithDuration(value time.Duration) { 335 | ts.timer.time(value) 336 | } 337 | 338 | type statStore struct { 339 | counters sync.Map 340 | gauges sync.Map 341 | timers sync.Map 342 | 343 | mu sync.RWMutex 344 | statGenerators []StatGenerator 345 | 346 | sink Sink 347 | } 348 | 349 | var ReservedTagWords = map[string]bool{"asg": true, "az": true, "backend": true, "canary": true, "host": true, "period": true, "region": true, "shard": true, "window": true, "source": true, "project": true, "facet": true, "envoyservice": true} 350 | 351 | func (s *statStore) validateTags(tags map[string]string) { 352 | for k := range tags { 353 | if _, ok := ReservedTagWords[k]; ok { 354 | // Keep track of how many times a reserved tag is used 355 | s.NewCounter("reserved_tag").Inc() 356 | } 357 | } 358 | } 359 | 360 | func (s *statStore) StartContext(ctx context.Context, ticker *time.Ticker) { 361 | for { 362 | select { 363 | case <-ctx.Done(): 364 | s.Flush() 365 | return 366 | case <-ticker.C: 367 | s.Flush() 368 | } 369 | } 370 | } 371 | 372 | func (s *statStore) Start(ticker *time.Ticker) { 373 | s.StartContext(context.Background(), ticker) 374 | } 375 | 376 | func (s *statStore) Flush() { 377 | s.mu.RLock() 378 | for _, g := range s.statGenerators { 379 | g.GenerateStats() 380 | } 381 | s.mu.RUnlock() 382 | 383 | s.counters.Range(func(key, v interface{}) bool { 384 | // do not flush counters that are set to zero 385 | if value := v.(*counter).latch(); value != 0 { 386 | s.sink.FlushCounter(key.(string), value) 387 | } 388 | return true 389 | }) 390 | 391 | s.gauges.Range(func(key, v interface{}) bool { 392 | s.sink.FlushGauge(key.(string), v.(*gauge).Value()) 393 | return true 394 | }) 395 | 396 | flushableSink, ok := s.sink.(FlushableSink) 397 | if ok { 398 | flushableSink.Flush() 399 | } 400 | } 401 | 402 | func (s *statStore) AddStatGenerator(statGenerator StatGenerator) { 403 | s.mu.Lock() 404 | defer s.mu.Unlock() 405 | s.statGenerators = append(s.statGenerators, statGenerator) 406 | } 407 | 408 | func (s *statStore) Store() Store { 409 | return s 410 | } 411 | 412 | func (s *statStore) Scope(name string) Scope { 413 | return newSubScope(s, name, nil) 414 | } 415 | 416 | func (s *statStore) ScopeWithTags(name string, tags map[string]string) Scope { 417 | s.validateTags(tags) 418 | return newSubScope(s, name, tags) 419 | } 420 | 421 | func (s *statStore) newCounter(serializedName string) *counter { 422 | if v, ok := s.counters.Load(serializedName); ok { 423 | return v.(*counter) 424 | } 425 | c := new(counter) 426 | if v, loaded := s.counters.LoadOrStore(serializedName, c); loaded { 427 | return v.(*counter) 428 | } 429 | return c 430 | } 431 | 432 | func (s *statStore) NewCounter(name string) Counter { 433 | return s.newCounter(name) 434 | } 435 | 436 | func (s *statStore) NewCounterWithTags(name string, tags map[string]string) Counter { 437 | s.validateTags(tags) 438 | return s.newCounter(tagspkg.SerializeTags(name, tags)) 439 | } 440 | 441 | func (s *statStore) newCounterWithTagSet(name string, tags tagspkg.TagSet) Counter { 442 | return s.newCounter(tags.Serialize(name)) 443 | } 444 | 445 | var emptyPerInstanceTags = map[string]string{"_f": "i"} 446 | 447 | func (s *statStore) NewPerInstanceCounter(name string, tags map[string]string) Counter { 448 | if len(tags) == 0 { 449 | return s.NewCounterWithTags(name, emptyPerInstanceTags) 450 | } 451 | if _, found := tags["_f"]; found { 452 | return s.NewCounterWithTags(name, tags) 453 | } 454 | s.validateTags(tags) 455 | return s.newCounterWithTagSet(name, tagspkg.TagSet(nil).MergePerInstanceTags(tags)) 456 | } 457 | 458 | func (s *statStore) newGauge(serializedName string) *gauge { 459 | if v, ok := s.gauges.Load(serializedName); ok { 460 | return v.(*gauge) 461 | } 462 | g := new(gauge) 463 | if v, loaded := s.gauges.LoadOrStore(serializedName, g); loaded { 464 | return v.(*gauge) 465 | } 466 | return g 467 | } 468 | 469 | func (s *statStore) NewGauge(name string) Gauge { 470 | return s.newGauge(name) 471 | } 472 | 473 | func (s *statStore) NewGaugeWithTags(name string, tags map[string]string) Gauge { 474 | s.validateTags(tags) 475 | return s.newGauge(tagspkg.SerializeTags(name, tags)) 476 | } 477 | 478 | func (s *statStore) newGaugeWithTagSet(name string, tags tagspkg.TagSet) Gauge { 479 | return s.newGauge(tags.Serialize(name)) 480 | } 481 | 482 | func (s *statStore) NewPerInstanceGauge(name string, tags map[string]string) Gauge { 483 | if len(tags) == 0 { 484 | return s.NewGaugeWithTags(name, emptyPerInstanceTags) 485 | } 486 | if _, found := tags["_f"]; found { 487 | return s.NewGaugeWithTags(name, tags) 488 | } 489 | s.validateTags(tags) 490 | return s.newGaugeWithTagSet(name, tagspkg.TagSet(nil).MergePerInstanceTags(tags)) 491 | } 492 | 493 | func (s *statStore) newTimer(serializedName string, base time.Duration) *timer { 494 | if v, ok := s.timers.Load(serializedName); ok { 495 | return v.(*timer) 496 | } 497 | t := &timer{name: serializedName, sink: s.sink, base: base} 498 | if v, loaded := s.timers.LoadOrStore(serializedName, t); loaded { 499 | return v.(*timer) 500 | } 501 | return t 502 | } 503 | 504 | func (s *statStore) NewMilliTimer(name string) Timer { 505 | return s.newTimer(name, time.Millisecond) 506 | } 507 | 508 | func (s *statStore) NewMilliTimerWithTags(name string, tags map[string]string) Timer { 509 | s.validateTags(tags) 510 | return s.newTimer(tagspkg.SerializeTags(name, tags), time.Millisecond) 511 | } 512 | 513 | func (s *statStore) NewTimer(name string) Timer { 514 | return s.newTimer(name, time.Microsecond) 515 | } 516 | 517 | func (s *statStore) NewTimerWithTags(name string, tags map[string]string) Timer { 518 | s.validateTags(tags) 519 | return s.newTimer(tagspkg.SerializeTags(name, tags), time.Microsecond) 520 | } 521 | 522 | func (s *statStore) newTimerWithTagSet(name string, tags tagspkg.TagSet, base time.Duration) Timer { 523 | return s.newTimer(tags.Serialize(name), base) 524 | } 525 | 526 | func (s *statStore) NewPerInstanceTimer(name string, tags map[string]string) Timer { 527 | if len(tags) == 0 { 528 | return s.NewTimerWithTags(name, emptyPerInstanceTags) 529 | } 530 | if _, found := tags["_f"]; found { 531 | return s.NewTimerWithTags(name, tags) 532 | } 533 | s.validateTags(tags) 534 | return s.newTimerWithTagSet(name, tagspkg.TagSet(nil).MergePerInstanceTags(tags), time.Microsecond) 535 | } 536 | 537 | func (s *statStore) NewPerInstanceMilliTimer(name string, tags map[string]string) Timer { 538 | if len(tags) == 0 { 539 | return s.NewMilliTimerWithTags(name, emptyPerInstanceTags) 540 | } 541 | if _, found := tags["_f"]; found { 542 | return s.NewMilliTimerWithTags(name, tags) 543 | } 544 | s.validateTags(tags) 545 | return s.newTimerWithTagSet(name, tagspkg.TagSet(nil).MergePerInstanceTags(tags), time.Millisecond) 546 | } 547 | 548 | type subScope struct { 549 | registry *statStore 550 | name string 551 | tags tagspkg.TagSet // read-only and may be shared by multiple subScopes 552 | } 553 | 554 | func newSubScope(registry *statStore, name string, tags map[string]string) *subScope { 555 | return &subScope{registry: registry, name: name, tags: tagspkg.NewTagSet(tags)} 556 | } 557 | 558 | func (s *subScope) Scope(name string) Scope { 559 | return s.ScopeWithTags(name, nil) 560 | } 561 | 562 | func (s *subScope) ScopeWithTags(name string, tags map[string]string) Scope { 563 | s.registry.validateTags(tags) 564 | return &subScope{ 565 | registry: s.registry, 566 | name: joinScopes(s.name, name), 567 | tags: s.tags.MergeTags(tags), 568 | } 569 | } 570 | 571 | func (s *subScope) Store() Store { 572 | return s.registry 573 | } 574 | 575 | func (s *subScope) NewCounter(name string) Counter { 576 | return s.NewCounterWithTags(name, nil) 577 | } 578 | 579 | func (s *subScope) NewCounterWithTags(name string, tags map[string]string) Counter { 580 | return s.registry.newCounterWithTagSet(joinScopes(s.name, name), s.tags.MergeTags(tags)) 581 | } 582 | 583 | func (s *subScope) NewPerInstanceCounter(name string, tags map[string]string) Counter { 584 | return s.registry.newCounterWithTagSet(joinScopes(s.name, name), 585 | s.tags.MergePerInstanceTags(tags)) 586 | } 587 | 588 | func (s *subScope) NewGauge(name string) Gauge { 589 | return s.NewGaugeWithTags(name, nil) 590 | } 591 | 592 | func (s *subScope) NewGaugeWithTags(name string, tags map[string]string) Gauge { 593 | return s.registry.newGaugeWithTagSet(joinScopes(s.name, name), s.tags.MergeTags(tags)) 594 | } 595 | 596 | func (s *subScope) NewPerInstanceGauge(name string, tags map[string]string) Gauge { 597 | return s.registry.newGaugeWithTagSet(joinScopes(s.name, name), 598 | s.tags.MergePerInstanceTags(tags)) 599 | } 600 | 601 | func (s *subScope) NewTimer(name string) Timer { 602 | return s.NewTimerWithTags(name, nil) 603 | } 604 | 605 | func (s *subScope) NewTimerWithTags(name string, tags map[string]string) Timer { 606 | return s.registry.newTimerWithTagSet(joinScopes(s.name, name), s.tags.MergeTags(tags), time.Microsecond) 607 | } 608 | 609 | func (s *subScope) NewPerInstanceTimer(name string, tags map[string]string) Timer { 610 | return s.registry.newTimerWithTagSet(joinScopes(s.name, name), 611 | s.tags.MergePerInstanceTags(tags), time.Microsecond) 612 | } 613 | 614 | func (s *subScope) NewMilliTimer(name string) Timer { 615 | return s.NewMilliTimerWithTags(name, nil) 616 | } 617 | 618 | func (s *subScope) NewMilliTimerWithTags(name string, tags map[string]string) Timer { 619 | return s.registry.newTimerWithTagSet(joinScopes(s.name, name), s.tags.MergeTags(tags), time.Millisecond) 620 | } 621 | 622 | func (s *subScope) NewPerInstanceMilliTimer(name string, tags map[string]string) Timer { 623 | s.registry.validateTags(tags) 624 | return s.registry.newTimerWithTagSet(joinScopes(s.name, name), 625 | s.tags.MergePerInstanceTags(tags), time.Millisecond) 626 | } 627 | 628 | func joinScopes(parent, child string) string { 629 | return parent + "." + child 630 | } 631 | -------------------------------------------------------------------------------- /net_sink_test.go: -------------------------------------------------------------------------------- 1 | package stats 2 | 3 | import ( 4 | "bufio" 5 | "context" 6 | "fmt" 7 | "io" 8 | "net" 9 | "os" 10 | "os/exec" 11 | "path/filepath" 12 | "runtime" 13 | "strconv" 14 | "strings" 15 | "sync" 16 | "sync/atomic" 17 | "testing" 18 | "time" 19 | ) 20 | 21 | func foreverNow() time.Time { 22 | return time.Date(2022, 1, 1, 0, 0, 0, 0, time.UTC) 23 | } 24 | 25 | type testStatSink struct { 26 | sync.Mutex 27 | record string 28 | } 29 | 30 | func (s *testStatSink) FlushCounter(name string, value uint64) { 31 | s.Lock() 32 | s.record += fmt.Sprintf("%s:%d|c\n", name, value) 33 | s.Unlock() 34 | } 35 | 36 | func (s *testStatSink) FlushGauge(name string, value uint64) { 37 | s.Lock() 38 | s.record += fmt.Sprintf("%s:%d|g\n", name, value) 39 | s.Unlock() 40 | } 41 | 42 | func (s *testStatSink) FlushTimer(name string, value float64) { 43 | s.Lock() 44 | s.record += fmt.Sprintf("%s:%f|ms\n", name, value) 45 | s.Unlock() 46 | } 47 | 48 | func TestCreateTimer(t *testing.T) { 49 | sink := &testStatSink{} 50 | store := NewStore(sink, true) 51 | 52 | t1 := store.NewTimer("hello") 53 | if t1 == nil { 54 | t.Fatal("No timer returned") 55 | } 56 | 57 | t2 := store.NewTimer("hello") 58 | if t1 != t2 { 59 | t.Error("A new timer with the same name was returned") 60 | } 61 | } 62 | 63 | func TestCreateTaggedTimer(t *testing.T) { 64 | sink := &testStatSink{} 65 | store := NewStore(sink, true) 66 | 67 | t1 := store.NewTimerWithTags("hello", map[string]string{"t1": "v1"}) 68 | if t1 == nil { 69 | t.Fatal("No timer returned") 70 | } 71 | 72 | t2 := store.NewTimerWithTags("hello", map[string]string{"t1": "v1"}) 73 | if t1 != t2 { 74 | t.Error("A new timer with the same name was returned") 75 | } 76 | } 77 | 78 | func TestCreateInstanceTimer(t *testing.T) { 79 | sink := &testStatSink{} 80 | store := NewStore(sink, true) 81 | 82 | t1 := store.NewPerInstanceTimer("hello", map[string]string{"t1": "v1"}) 83 | if t1 == nil { 84 | t.Fatal("No timer returned") 85 | } 86 | 87 | t2 := store.NewPerInstanceTimer("hello", map[string]string{"t1": "v1"}) 88 | if t1 != t2 { 89 | t.Error("A new timer with the same name was returned") 90 | } 91 | 92 | span := t2.AllocateSpan() 93 | span.Complete() 94 | 95 | expected := "hello.___f=i.__t1=v1" 96 | if !strings.Contains(sink.record, expected) { 97 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 98 | } 99 | } 100 | 101 | func TestCreateInstanceTimerNilMap(t *testing.T) { 102 | sink := &testStatSink{} 103 | store := NewStore(sink, true) 104 | 105 | t1 := store.NewPerInstanceTimer("hello", nil) 106 | if t1 == nil { 107 | t.Fatal("No timer returned") 108 | } 109 | 110 | t2 := store.NewPerInstanceTimer("hello", map[string]string{"_f": "i"}) 111 | if t1 != t2 { 112 | t.Error("A new timer with the same name was returned") 113 | } 114 | 115 | span := t2.AllocateSpan() 116 | span.Complete() 117 | 118 | expected := "hello.___f=i" 119 | if !strings.Contains(sink.record, expected) { 120 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 121 | } 122 | } 123 | 124 | func TestCreateDuplicateCounter(t *testing.T) { 125 | sink := &testStatSink{} 126 | store := NewStore(sink, true) 127 | 128 | t1 := store.NewCounter("TestCreateCounter") 129 | if t1 == nil { 130 | t.Fatal("No counter returned") 131 | } 132 | 133 | t2 := store.NewCounter("TestCreateCounter") 134 | if t1 != t2 { 135 | t.Error("A new counter with the same name was returned") 136 | } 137 | } 138 | 139 | func TestCounter(t *testing.T) { 140 | sink := &testStatSink{} 141 | store := NewStore(sink, true) 142 | 143 | counter := store.NewCounter("c") 144 | counter.Inc() 145 | counter.Add(10) 146 | counter.Inc() 147 | store.Flush() 148 | 149 | expected := "c:12|c\n" 150 | if expected != sink.record { 151 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 152 | } 153 | } 154 | 155 | func TestTaggedPerInstanceCounter(t *testing.T) { 156 | sink := &testStatSink{} 157 | store := NewStore(sink, true) 158 | scope := store.Scope("prefix") 159 | 160 | counter := scope.NewPerInstanceCounter("c", map[string]string{"tag1": "value1"}) 161 | counter.Inc() 162 | counter.Add(10) 163 | counter.Inc() 164 | store.Flush() 165 | 166 | expected := "prefix.c.___f=i.__tag1=value1:12|c\n" 167 | if expected != sink.record { 168 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 169 | } 170 | } 171 | 172 | func TestTaggedPerInstanceCounterWithNilTags(t *testing.T) { 173 | sink := &testStatSink{} 174 | store := NewStore(sink, true) 175 | scope := store.Scope("prefix") 176 | 177 | counter := scope.NewPerInstanceCounter("c", nil) 178 | counter.Inc() 179 | counter.Add(10) 180 | counter.Inc() 181 | store.Flush() 182 | 183 | expected := "prefix.c.___f=i:12|c\n" 184 | if expected != sink.record { 185 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 186 | } 187 | } 188 | 189 | func TestTaggedCounter(t *testing.T) { 190 | sink := &testStatSink{} 191 | store := NewStore(sink, true) 192 | scope := store.Scope("prefix") 193 | 194 | counter := scope.NewCounterWithTags("c", map[string]string{"tag1": "value1"}) 195 | counter.Inc() 196 | counter.Add(10) 197 | counter.Inc() 198 | store.Flush() 199 | 200 | expected := "prefix.c.__tag1=value1:12|c\n" 201 | if expected != sink.record { 202 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 203 | } 204 | } 205 | 206 | func TestCreateDuplicateGauge(t *testing.T) { 207 | sink := &testStatSink{} 208 | store := NewStore(sink, true) 209 | 210 | t1 := store.NewGauge("TestCreateGauge") 211 | if t1 == nil { 212 | t.Fatal("No gauge returned") 213 | } 214 | 215 | t2 := store.NewGauge("TestCreateGauge") 216 | if t1 != t2 { 217 | t.Error("A new counter with the same name was returned") 218 | } 219 | } 220 | 221 | func TestGauge(t *testing.T) { 222 | sink := &testStatSink{} 223 | store := NewStore(sink, true) 224 | 225 | gauge := store.NewGauge("TestGauge") 226 | gauge.Inc() 227 | gauge.Add(10) 228 | gauge.Dec() 229 | gauge.Sub(5) 230 | store.Flush() 231 | 232 | expected := "TestGauge:5|g\n" 233 | if expected != sink.record { 234 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 235 | } 236 | } 237 | 238 | func TestTaggedGauge(t *testing.T) { 239 | sink := &testStatSink{} 240 | store := NewStore(sink, true) 241 | 242 | gauge := store.NewGaugeWithTags("TestGauge", map[string]string{"tag1": "v1"}) 243 | gauge.Inc() 244 | gauge.Add(10) 245 | gauge.Dec() 246 | gauge.Sub(5) 247 | store.Flush() 248 | 249 | expected := "TestGauge.__tag1=v1:5|g\n" 250 | if expected != sink.record { 251 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 252 | } 253 | } 254 | 255 | func TestPerInstanceGauge(t *testing.T) { 256 | sink := &testStatSink{} 257 | store := NewStore(sink, true) 258 | 259 | gauge := store.NewPerInstanceGauge("TestGauge", map[string]string{"tag1": "v1"}) 260 | gauge.Inc() 261 | gauge.Add(10) 262 | gauge.Dec() 263 | gauge.Sub(5) 264 | store.Flush() 265 | 266 | expected := "TestGauge.___f=i.__tag1=v1:5|g\n" 267 | if expected != sink.record { 268 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 269 | } 270 | } 271 | 272 | func TestPerInstanceGaugeWithNilTags(t *testing.T) { 273 | sink := &testStatSink{} 274 | store := NewStore(sink, true) 275 | 276 | gauge := store.NewPerInstanceGauge("TestGauge", nil) 277 | gauge.Inc() 278 | gauge.Add(10) 279 | gauge.Dec() 280 | gauge.Sub(5) 281 | store.Flush() 282 | 283 | expected := "TestGauge.___f=i:5|g\n" 284 | if expected != sink.record { 285 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 286 | } 287 | } 288 | 289 | func TestScopes(t *testing.T) { 290 | sink := &testStatSink{} 291 | store := NewStore(sink, true) 292 | 293 | ascope := store.Scope("a") 294 | bscope := ascope.Scope("b") 295 | counter := bscope.NewCounter("c") 296 | counter.Inc() 297 | store.Flush() 298 | 299 | expected := "a.b.c:1|c\n" 300 | if expected != sink.record { 301 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 302 | } 303 | } 304 | 305 | func TestScopesWithTags(t *testing.T) { 306 | sink := &testStatSink{} 307 | store := NewStore(sink, true) 308 | 309 | ascope := store.ScopeWithTags("a", map[string]string{"x": "a", "y": "a"}) 310 | bscope := ascope.ScopeWithTags("b", map[string]string{"x": "b", "z": "b"}) 311 | dscope := bscope.Scope("d") 312 | counter := dscope.NewCounter("c") 313 | counter.Inc() 314 | timer := dscope.NewTimer("t") 315 | timer.AddValue(1) 316 | gauge := dscope.NewGauge("g") 317 | gauge.Set(1) 318 | store.Flush() 319 | 320 | expected := "a.b.d.t.__x=b.__y=a.__z=b:1.000000|ms\na.b.d.c.__x=b.__y=a.__z=b:1|c\na.b.d.g.__x=b.__y=a.__z=b:1|g\n" 321 | if expected != sink.record { 322 | t.Errorf("\n# Expected:\n%s\n# Got:\n%s\n", expected, sink.record) 323 | } 324 | } 325 | 326 | func TestScopesAndMetricsWithTags(t *testing.T) { 327 | sink := &testStatSink{} 328 | store := NewStore(sink, true) 329 | 330 | ascope := store.ScopeWithTags("a", map[string]string{"x": "a", "y": "a"}) 331 | bscope := ascope.Scope("b") 332 | counter := bscope.NewCounterWithTags("c", map[string]string{"x": "m", "z": "m"}) 333 | counter.Inc() 334 | timer := bscope.NewTimerWithTags("t", map[string]string{"x": "m", "z": "m"}) 335 | timer.AddValue(1) 336 | gauge := bscope.NewGaugeWithTags("g", map[string]string{"x": "m", "z": "m"}) 337 | gauge.Set(1) 338 | store.Flush() 339 | 340 | expected := "a.b.t.__x=m.__y=a.__z=m:1.000000|ms\na.b.c.__x=m.__y=a.__z=m:1|c\na.b.g.__x=m.__y=a.__z=m:1|g\n" 341 | if expected != sink.record { 342 | t.Errorf("\n# Expected:\n%s\n# Got:\n%s\n", expected, sink.record) 343 | } 344 | } 345 | 346 | type testStatGenerator struct { 347 | counter Counter 348 | gauge Gauge 349 | } 350 | 351 | func (s *testStatGenerator) GenerateStats() { 352 | s.counter.Add(123) 353 | s.gauge.Set(456) 354 | } 355 | 356 | func TestStatGenerator(t *testing.T) { 357 | sink := &testStatSink{} 358 | store := NewStore(sink, true) 359 | scope := store.Scope("TestRuntime") 360 | 361 | g := testStatGenerator{counter: scope.NewCounter("counter"), gauge: scope.NewGauge("gauge")} 362 | 363 | store.AddStatGenerator(&g) 364 | store.Flush() 365 | 366 | expected := "TestRuntime.counter:123|c\nTestRuntime.gauge:456|g\n" 367 | if expected != sink.record { 368 | t.Errorf("Expected: '%s' Got: '%s'", expected, sink.record) 369 | } 370 | } 371 | 372 | func TestNetSink_Flush(t *testing.T) { 373 | lc, err := net.Listen("tcp", "127.0.0.1:0") 374 | if err != nil { 375 | t.Fatal(err) 376 | } 377 | defer lc.Close() 378 | 379 | _, port, err := net.SplitHostPort(lc.Addr().String()) 380 | if err != nil { 381 | t.Fatal(err) 382 | } 383 | 384 | go func() { 385 | for { 386 | conn, err := lc.Accept() 387 | if err == nil { 388 | _, err = io.Copy(io.Discard, conn) 389 | } 390 | if err != nil { 391 | return 392 | } 393 | } 394 | }() 395 | 396 | nport, err := strconv.Atoi(port) 397 | if err != nil { 398 | t.Fatal(err) 399 | } 400 | sink := NewNetSink( 401 | WithLogger(discardLogger()), 402 | WithStatsdPort(nport), 403 | ) 404 | 405 | // Spin up a goroutine to flood the sink with large Counter stats. 406 | // The goal here is to keep the buffer channel full. 407 | ready := make(chan struct{}, 1) 408 | done := make(chan struct{}) 409 | defer close(done) 410 | go func() { 411 | name := "test." + strings.Repeat("a", 1024) + ".counter" 412 | for { 413 | select { 414 | case ready <- struct{}{}: 415 | // signal that we've started 416 | case <-done: 417 | return 418 | default: 419 | for i := 0; i < 1000; i++ { 420 | sink.FlushCounter(name, 1) 421 | } 422 | } 423 | } 424 | }() 425 | <-ready // wait for goroutine to start 426 | 427 | t.Run("One", func(t *testing.T) { 428 | flushed := make(chan struct{}) 429 | go func() { 430 | defer close(flushed) 431 | sink.Flush() 432 | }() 433 | select { 434 | case <-flushed: 435 | // ok 436 | case <-time.After(time.Second): 437 | t.Fatal("Flush blocked") 438 | } 439 | }) 440 | 441 | t.Run("Ten", func(t *testing.T) { 442 | flushed := make(chan struct{}) 443 | go func() { 444 | defer close(flushed) 445 | for i := 0; i < 10; i++ { 446 | sink.Flush() 447 | } 448 | }() 449 | select { 450 | case <-flushed: 451 | // ok 452 | case <-time.After(time.Second): 453 | t.Fatal("Flush blocked") 454 | } 455 | }) 456 | 457 | t.Run("Parallel", func(t *testing.T) { 458 | start := make(chan struct{}) 459 | wg := new(sync.WaitGroup) 460 | for i := 0; i < 20; i++ { 461 | wg.Add(1) 462 | go func() { 463 | <-start 464 | defer wg.Done() 465 | for i := 0; i < 10; i++ { 466 | sink.Flush() 467 | } 468 | }() 469 | } 470 | flushed := make(chan struct{}) 471 | go func() { 472 | close(start) 473 | wg.Wait() 474 | close(flushed) 475 | }() 476 | select { 477 | case <-flushed: 478 | // ok 479 | case <-time.After(time.Second): 480 | t.Fatal("Flush blocked") 481 | } 482 | }) 483 | } 484 | 485 | // Test that drainFlushQueue() does not hang when there are continuous 486 | // flush requests. 487 | func TestNetSink_DrainFlushQueue(t *testing.T) { 488 | s := &netSink{ 489 | doFlush: make(chan chan struct{}, 8), 490 | } 491 | 492 | sent := new(int64) 493 | 494 | // Saturate the flush channel 495 | 496 | done := make(chan struct{}) 497 | defer close(done) 498 | 499 | for i := 0; i < runtime.NumCPU(); i++ { 500 | go func() { 501 | for { 502 | select { 503 | case s.doFlush <- make(chan struct{}): 504 | atomic.AddInt64(sent, 1) 505 | case <-done: 506 | return 507 | } 508 | } 509 | }() 510 | } 511 | 512 | // Wait for the flush channel to fill 513 | for len(s.doFlush) < cap(s.doFlush) { 514 | runtime.Gosched() 515 | } 516 | 517 | flushed := make(chan struct{}) 518 | go func() { 519 | s.drainFlushQueue() 520 | close(flushed) 521 | }() 522 | 523 | // We will flush up to cap(s.doFlush) * 8 items, so the max number 524 | // of sends will be that plus the capacity of the buffer. 525 | maxSends := cap(s.doFlush)*8 + cap(s.doFlush) 526 | 527 | select { 528 | case <-flushed: 529 | n := int(atomic.LoadInt64(sent)) 530 | switch { 531 | case n < cap(s.doFlush): 532 | // This should be impossible since we fill the channel 533 | // before calling drainFlushQueue(). 534 | t.Errorf("Sent less than %d items: %d", cap(s.doFlush), n) 535 | case n > maxSends: 536 | // This should be nearly impossible to get without inserting 537 | // runtime.Gosched() into the flush/drain loop. 538 | t.Errorf("Sent more than %d items: %d", maxSends, n) 539 | } 540 | case <-time.After(time.Second / 2): 541 | // 500ms is really generous, it should return almost immediately. 542 | t.Error("drainFlushQueue did not return in time") 543 | } 544 | } 545 | 546 | func discardLogger() *loggingSink { 547 | return &loggingSink{writer: io.Discard, now: foreverNow} 548 | } 549 | 550 | func setupTestNetSink(t *testing.T, protocol string, stop bool) (*netTestSink, *netSink) { 551 | ts := newNetTestSink(t, protocol) 552 | 553 | if stop { 554 | if err := ts.Close(); err != nil { 555 | t.Fatal(err) 556 | } 557 | } 558 | 559 | sink := NewTCPStatsdSink( 560 | WithLogger(discardLogger()), 561 | WithStatsdHost(ts.Host(t)), 562 | WithStatsdPort(ts.Port(t)), 563 | WithStatsdProtocol(protocol), 564 | ).(*netSink) 565 | 566 | return ts, sink 567 | } 568 | 569 | func testNetSinkBufferSize(t *testing.T, protocol string) { 570 | var size int 571 | switch protocol { 572 | case "udp": 573 | size = defaultBufferSizeUDP 574 | case "tcp": 575 | size = defaultBufferSizeTCP 576 | } 577 | ts, sink := setupTestNetSink(t, protocol, false) 578 | defer ts.Close() 579 | if sink.bufWriter.Size() != size { 580 | t.Errorf("Buffer Size: got: %d want: %d", sink.bufWriter.Size(), size) 581 | } 582 | } 583 | 584 | func testNetSinkStatTypes(t *testing.T, protocol string) { 585 | expected := [...]string{ 586 | "counter:1|c\n", 587 | "gauge:1|g\n", 588 | "timer_int:1|ms\n", 589 | "timer_float:1.230000|ms\n", 590 | } 591 | 592 | ts, sink := setupTestNetSink(t, protocol, false) 593 | defer ts.Close() 594 | 595 | sink.FlushCounter("counter", 1) 596 | sink.FlushGauge("gauge", 1) 597 | sink.FlushTimer("timer_int", 1) 598 | sink.FlushTimer("timer_float", 1.23) 599 | sink.Flush() 600 | 601 | for _, exp := range expected { 602 | stat := ts.WaitForStat(t, time.Millisecond*50) 603 | if stat != exp { 604 | t.Errorf("stats got: %q want: %q", stat, exp) 605 | } 606 | } 607 | 608 | // make sure there aren't any extra stats we're missing 609 | exp := strings.Join(expected[:], "") 610 | buf := ts.String() 611 | if buf != exp { 612 | t.Errorf("stats buffer\ngot:\n%q\nwant:\n%q\n", buf, exp) 613 | } 614 | } 615 | 616 | func testNetSinkImmediateFlush(t *testing.T, protocol string) { 617 | const expected = "counter:1|c\n" 618 | 619 | ts, sink := setupTestNetSink(t, protocol, false) 620 | defer ts.Close() 621 | 622 | sink.FlushCounter("counter", 1) 623 | sink.Flush() 624 | 625 | stat := ts.WaitForStat(t, time.Millisecond*50) 626 | if stat != expected { 627 | t.Errorf("stats got: %q want: %q", stat, expected) 628 | } 629 | } 630 | 631 | // replaceFatalWithLog replaces calls to t.Fatalf() with t.Logf(). 632 | type replaceFatalWithLog struct { 633 | *testing.T 634 | } 635 | 636 | func (t replaceFatalWithLog) Fatalf(format string, args ...interface{}) { 637 | t.Logf(format, args...) 638 | } 639 | 640 | func testNetSinkReconnect(t *testing.T, protocol string) { 641 | if testing.Short() { 642 | t.Skip("Skipping: short test") 643 | } 644 | t.Parallel() 645 | 646 | const expected = "counter:1|c\n" 647 | 648 | ts, sink := setupTestNetSink(t, protocol, true) 649 | defer ts.Close() 650 | 651 | sink.FlushCounter("counter", 1) 652 | 653 | flushed := make(chan struct{}) 654 | go func() { 655 | flushed <- struct{}{} 656 | sink.Flush() 657 | close(flushed) 658 | }() 659 | 660 | <-flushed // wait till we're ready 661 | ts.Restart(t, true) 662 | 663 | sink.FlushCounter("counter", 1) 664 | 665 | // This test is flaky with UDP and the race detector, but good 666 | // to have so we log instead of fail the test. 667 | if protocol == "udp" { 668 | stat := ts.WaitForStat(replaceFatalWithLog{t}, defaultRetryInterval*3) 669 | if stat != "" && stat != expected { 670 | t.Fatalf("stats got: %q want: %q", stat, expected) 671 | } 672 | } else { 673 | stat := ts.WaitForStat(t, defaultRetryInterval*3) 674 | if stat != expected { 675 | t.Fatalf("stats got: %q want: %q", stat, expected) 676 | } 677 | } 678 | 679 | // Make sure our flush call returned 680 | select { 681 | case <-flushed: 682 | case <-time.After(time.Millisecond * 100): 683 | // The flushed channel should be closed by this point, 684 | // but this was failing in CI on go1.12 due to timing 685 | // issues so we relax the constraint and give it 100ms. 686 | t.Error("Flush() did not return") 687 | } 688 | } 689 | 690 | func testNetSinkReconnectFailure(t *testing.T, protocol string) { 691 | if testing.Short() { 692 | t.Skip("Skipping: short test") 693 | } 694 | t.Parallel() 695 | 696 | ts, sink := setupTestNetSink(t, protocol, true) 697 | defer ts.Close() 698 | 699 | sink.FlushCounter("counter", 1) 700 | 701 | const N = 16 702 | flushCount := new(int64) 703 | flushed := make(chan struct{}) 704 | go func() { 705 | wg := new(sync.WaitGroup) 706 | wg.Add(N) 707 | for i := 0; i < N; i++ { 708 | go func() { 709 | sink.Flush() 710 | atomic.AddInt64(flushCount, 1) 711 | wg.Done() 712 | }() 713 | } 714 | wg.Wait() 715 | close(flushed) 716 | }() 717 | 718 | // Make sure our flush call returned 719 | select { 720 | case <-flushed: 721 | // Ok 722 | case <-time.After(defaultRetryInterval * 2): 723 | t.Fatalf("Only %d of %d Flush() calls succeeded", 724 | atomic.LoadInt64(flushCount), N) 725 | } 726 | } 727 | 728 | func TestNetSink_BufferSize_TCP(t *testing.T) { 729 | testNetSinkBufferSize(t, "tcp") 730 | } 731 | 732 | func TestNetSink_BufferSize_UDP(t *testing.T) { 733 | testNetSinkBufferSize(t, "udp") 734 | } 735 | 736 | func TestNetSink_StatTypes_TCP(t *testing.T) { 737 | testNetSinkStatTypes(t, "tcp") 738 | } 739 | 740 | func TestNetSink_StatTypes_UDP(t *testing.T) { 741 | testNetSinkStatTypes(t, "udp") 742 | } 743 | 744 | func TestNetSink_ImmediateFlush_TCP(t *testing.T) { 745 | testNetSinkImmediateFlush(t, "tcp") 746 | } 747 | 748 | func TestNetSink_ImmediateFlush_UDP(t *testing.T) { 749 | testNetSinkImmediateFlush(t, "udp") 750 | } 751 | 752 | func TestNetSink_Reconnect_TCP(t *testing.T) { 753 | testNetSinkReconnect(t, "tcp") 754 | } 755 | 756 | func TestNetSink_Reconnect_UDP(t *testing.T) { 757 | testNetSinkReconnect(t, "udp") 758 | } 759 | 760 | func TestNetSink_ReconnectFailure_TCP(t *testing.T) { 761 | testNetSinkReconnectFailure(t, "tcp") 762 | } 763 | 764 | func TestNetSink_ReconnectFailure_UDP(t *testing.T) { 765 | testNetSinkReconnectFailure(t, "udp") 766 | } 767 | 768 | func buildBinary(t testing.TB, path string) (string, func()) { 769 | var binaryName string 770 | if strings.HasSuffix(path, ".go") { 771 | // foo/bar/main.go => bar 772 | binaryName = filepath.Base(filepath.Dir(path)) 773 | } else { 774 | binaryName = filepath.Base(path) 775 | } 776 | 777 | tmpdir, err := os.MkdirTemp("", "gostats-") 778 | if err != nil { 779 | t.Fatalf("creating tempdir: %v", err) 780 | } 781 | output := filepath.Join(tmpdir, binaryName) 782 | 783 | out, err := exec.Command("go", "build", "-o", output, path).CombinedOutput() 784 | if err != nil { 785 | t.Fatalf("failed to build %s: %s\n### output:\n%s\n###\n", 786 | path, err, strings.TrimSpace(string(out))) 787 | } 788 | 789 | cleanup := func() { 790 | os.RemoveAll(tmpdir) 791 | } 792 | return output, cleanup 793 | } 794 | 795 | func testNetSinkIntegration(t *testing.T, protocol string) { 796 | t.Parallel() 797 | 798 | ctx, cancel := context.WithCancel(context.Background()) 799 | defer cancel() 800 | 801 | fastExitExe, deleteBinary := buildBinary(t, "testdata/fast_exit/fast_exit.go") 802 | defer deleteBinary() 803 | 804 | // Test the stats of a fast exiting program are captured. 805 | t.Run("FastExit", func(t *testing.T) { 806 | ts := newNetTestSink(t, protocol) 807 | defer ts.Close() 808 | 809 | cmd := exec.CommandContext(ctx, fastExitExe) 810 | cmd.Env = ts.CommandEnv(t) 811 | 812 | out, err := cmd.CombinedOutput() 813 | if err != nil { 814 | t.Fatalf("Running command: %s\n### output:\n%s\n###\n", 815 | fastExitExe, strings.TrimSpace(string(out))) 816 | } 817 | 818 | stats := ts.String() 819 | const expected = "test.fast.exit.counter:1|c\n" 820 | if stats != expected { 821 | t.Errorf("stats: got: %q want: %q", stats, expected) 822 | } 823 | }) 824 | 825 | // Test that Flush() does not hang if the TCP sink is in a reconnect loop 826 | t.Run("Reconnect", func(t *testing.T) { 827 | ts := newNetTestSink(t, protocol) 828 | defer ts.Close() 829 | 830 | cmd := exec.CommandContext(ctx, fastExitExe) 831 | cmd.Env = ts.CommandEnv(t) 832 | 833 | if err := cmd.Start(); err != nil { 834 | t.Fatal(err) 835 | } 836 | errCh := make(chan error, 1) 837 | go func() { errCh <- cmd.Wait() }() 838 | 839 | select { 840 | case err := <-errCh: 841 | if err != nil { 842 | t.Fatal(err) 843 | } 844 | case <-time.After(defaultRetryInterval * 2): 845 | t.Fatal("Timed out waiting for command to exit") 846 | } 847 | }) 848 | } 849 | 850 | func TestNetSink_Integration_TCP(t *testing.T) { 851 | testNetSinkIntegration(t, "tcp") 852 | } 853 | 854 | func TestNetSink_Integration_UDP(t *testing.T) { 855 | testNetSinkIntegration(t, "udp") 856 | } 857 | 858 | type nopWriter struct{} 859 | 860 | func (nopWriter) Write(b []byte) (int, error) { 861 | return len(b), nil 862 | } 863 | 864 | func BenchmarkFlushCounter(b *testing.B) { 865 | sink := netSink{ 866 | bufWriter: bufio.NewWriter(nopWriter{}), 867 | } 868 | for i := 0; i < b.N; i++ { 869 | sink.FlushCounter("TestCounter.___f=i.__tag1=v1", uint64(i)) 870 | } 871 | } 872 | 873 | func BenchmarkFlushTimer(b *testing.B) { 874 | sink := netSink{ 875 | bufWriter: bufio.NewWriter(nopWriter{}), 876 | } 877 | for i := 0; i < b.N; i++ { 878 | sink.FlushTimer("TestTImer.___f=i.__tag1=v1", float64(i)/3) 879 | } 880 | } 881 | -------------------------------------------------------------------------------- /internal/tags/tags_test.go: -------------------------------------------------------------------------------- 1 | package tags 2 | 3 | import ( 4 | "bytes" 5 | crand "crypto/rand" 6 | "encoding/hex" 7 | "fmt" 8 | "math/rand" 9 | "reflect" 10 | "runtime" 11 | "sort" 12 | "strconv" 13 | "sync" 14 | "testing" 15 | "time" 16 | "unsafe" 17 | ) 18 | 19 | // Reference serializeTags implementation 20 | func serializeTagsReference(name string, tags map[string]string) string { 21 | const prefix = ".__" 22 | const sep = "=" 23 | if len(tags) == 0 { 24 | return name 25 | } 26 | tagPairs := make([]Tag, 0, len(tags)) 27 | for tagKey, tagValue := range tags { 28 | tagValue = ReplaceChars(tagValue) 29 | tagPairs = append(tagPairs, Tag{tagKey, tagValue}) 30 | } 31 | sort.Sort(TagSet(tagPairs)) 32 | 33 | buf := new(bytes.Buffer) 34 | for _, tag := range tagPairs { 35 | if tag.Key != "" && tag.Value != "" { 36 | fmt.Fprint(buf, prefix, tag.Key, sep, tag.Value) 37 | } 38 | } 39 | return name + buf.String() 40 | } 41 | 42 | func TestSerializeTags(t *testing.T) { 43 | const name = "prefix" 44 | const expected = name + ".__q=r.__zzz=hello" 45 | tags := map[string]string{"zzz": "hello", "q": "r"} 46 | serialized := SerializeTags(name, tags) 47 | if serialized != expected { 48 | t.Errorf("Serialized output (%s) didn't match expected output: %s", 49 | serialized, expected) 50 | } 51 | } 52 | 53 | // Test that the optimized serializeTags() function matches the reference 54 | // implementation. 55 | func TestSerializeTagsReference(t *testing.T) { 56 | const name = "prefix" 57 | makeTags := func(n int) map[string]string { 58 | m := make(map[string]string, n) 59 | for i := 0; i < n; i++ { 60 | k := fmt.Sprintf("key%d", i) 61 | v := fmt.Sprintf("val%d", i) 62 | m[k] = v 63 | } 64 | return m 65 | } 66 | for i := 0; i < 100; i++ { 67 | tags := makeTags(i) 68 | expected := serializeTagsReference(name, tags) 69 | serialized := SerializeTags(name, tags) 70 | if serialized != expected { 71 | t.Errorf("%d Serialized output (%s) didn't match expected output: %s", 72 | i, serialized, expected) 73 | } 74 | } 75 | } 76 | 77 | // Test the network sort used when we have 4 or less tags. Since the iteration 78 | // order of maps is random we use random keys in an attempt to get 100% test 79 | // coverage. 80 | func TestSerializeTagsNetworkSort(t *testing.T) { 81 | const name = "prefix" 82 | 83 | rand.Seed(time.Now().UnixNano()) 84 | seen := make(map[string]bool) 85 | 86 | randString := func() string { 87 | for i := 0; i < 100; i++ { 88 | s := randomString(t, rand.Intn(30)+1) 89 | if !seen[s] { 90 | seen[s] = true 91 | return s 92 | } 93 | } 94 | t.Fatal("Failed to generate a random string") 95 | return "" 96 | } 97 | 98 | makeTags := func(n int) map[string]string { 99 | m := make(map[string]string, n) 100 | for i := 0; i < n; i++ { 101 | k := randString() 102 | v := randString() 103 | m[k] = v 104 | } 105 | return m 106 | } 107 | 108 | // we use a network sort when tag length is 4 or less, but test up to 8 109 | // here in case that value is ever increased. 110 | for i := 1; i <= 4; i++ { 111 | // loop to increase the odds of 100% test coverage 112 | for i := 0; i < 10; i++ { 113 | tags := makeTags(i) 114 | expected := serializeTagsReference(name, tags) 115 | serialized := SerializeTags(name, tags) 116 | if serialized != expected { 117 | t.Errorf("%d Serialized output (%s) didn't match expected output: %s", 118 | i, serialized, expected) 119 | } 120 | } 121 | } 122 | } 123 | 124 | func TestSerializeTagsInvalidKeyValue(t *testing.T) { 125 | // Baseline tests against a hardcoded expected value 126 | t.Run("Baseline", func(t *testing.T) { 127 | const expected = "name.__1=1" 128 | tags := map[string]string{ 129 | "": "invalid_key", 130 | "invalid_value": "", 131 | "1": "1", 132 | } 133 | orig := make(map[string]string) 134 | for k, v := range tags { 135 | orig[k] = v 136 | } 137 | 138 | s := SerializeTags("name", tags) 139 | if s != expected { 140 | t.Errorf("Serialized output (%s) didn't match expected output: %s", 141 | s, expected) 142 | } 143 | 144 | if !reflect.DeepEqual(tags, orig) { 145 | t.Errorf("serializeTags modified the input map: %+v want: %+v", tags, orig) 146 | } 147 | }) 148 | 149 | createTags := func(n int) map[string]string { 150 | tags := make(map[string]string) 151 | for i := 0; i < n; i++ { 152 | key := fmt.Sprintf("key_%d", i) 153 | val := fmt.Sprintf("val_%d", i) 154 | tags[key] = val 155 | } 156 | return tags 157 | } 158 | 159 | test := func(t *testing.T, tags map[string]string) { 160 | orig := make(map[string]string) 161 | for k, v := range tags { 162 | orig[k] = v 163 | } 164 | 165 | got := SerializeTags("name", tags) 166 | exp := serializeTagsReference("name", tags) 167 | if got != exp { 168 | t.Errorf("Tags (%d) got: %q want: %q", len(tags), got, exp) 169 | } 170 | 171 | if !reflect.DeepEqual(tags, orig) { 172 | t.Errorf("serializeTags modified the input map: %+v want: %+v", tags, orig) 173 | } 174 | } 175 | 176 | t.Run("EmptyValue", func(t *testing.T) { 177 | for n := 0; n <= 10; n++ { 178 | tags := createTags(n) 179 | tags["invalid"] = "" 180 | test(t, tags) 181 | } 182 | }) 183 | 184 | t.Run("EmptyKey", func(t *testing.T) { 185 | for n := 0; n <= 10; n++ { 186 | tags := createTags(n) 187 | tags[""] = "invalid" 188 | test(t, tags) 189 | } 190 | }) 191 | 192 | t.Run("EmptyKeyValue", func(t *testing.T) { 193 | for n := 0; n <= 10; n++ { 194 | tags := createTags(n) 195 | tags[""] = "invalid" 196 | tags["invalid"] = "" 197 | test(t, tags) 198 | } 199 | }) 200 | } 201 | 202 | func TestSerializeTagsInvalidKeyValue_ThreadSafe(t *testing.T) { 203 | tags := map[string]string{ 204 | "": "invalid_key", 205 | "1": "1", 206 | } 207 | // Add some more keys to slow this down 208 | for i := 0; i < 256; i++ { 209 | v := "val_" + strconv.Itoa(i) 210 | k := "key_" + v 211 | tags[k] = v 212 | } 213 | 214 | // Make a copy 215 | orig := make(map[string]string, len(tags)) 216 | for k, v := range tags { 217 | orig[k] = v 218 | } 219 | 220 | start := make(chan struct{}) 221 | var wg sync.WaitGroup 222 | for i := 0; i < runtime.NumCPU()*2; i++ { 223 | wg.Add(1) 224 | go func() { 225 | defer wg.Done() 226 | <-start 227 | SerializeTags("name", tags) 228 | }() 229 | } 230 | close(start) 231 | wg.Wait() 232 | 233 | if !reflect.DeepEqual(tags, orig) { 234 | t.Error("serializeTags modified the input map") 235 | } 236 | } 237 | 238 | func TestSerializeWithPerInstanceFlag(t *testing.T) { 239 | const name = "prefix" 240 | const expected = name + ".___f=i.__foo=bar" 241 | tags := map[string]string{"foo": "bar", "_f": "i"} 242 | serialized := SerializeTags(name, tags) 243 | if serialized != expected { 244 | t.Errorf("Serialized output (%s) didn't match expected output: %s", 245 | serialized, expected) 246 | } 247 | } 248 | 249 | func TestSerializeIllegalTags(t *testing.T) { 250 | const name = "prefix" 251 | const expected = name + ".__foo=b_a_r.__q=p" 252 | tags := map[string]string{"foo": "b|a:r", "q": "p"} 253 | serialized := SerializeTags(name, tags) 254 | if serialized != expected { 255 | t.Errorf("Serialized output (%s) didn't match expected output: %s", 256 | serialized, expected) 257 | } 258 | } 259 | 260 | func TestSerializeTagValuePeriod(t *testing.T) { 261 | const name = "prefix" 262 | const expected = name + ".__foo=blah_blah.__q=p" 263 | tags := map[string]string{"foo": "blah.blah", "q": "p"} 264 | serialized := SerializeTags(name, tags) 265 | if serialized != expected { 266 | t.Errorf("Serialized output (%s) didn't match expected output: %s", 267 | serialized, expected) 268 | } 269 | } 270 | 271 | func TestSerializeTagDiscardEmptyTagKeyValue(t *testing.T) { 272 | const name = "prefix" 273 | const expected = name + ".__key1=value1.__key3=value3" 274 | tags := map[string]string{"key1": "value1", "key2": "", "key3": "value3", "": "value4"} 275 | serialized := SerializeTags(name, tags) 276 | if serialized != expected { 277 | t.Errorf("Serialized output (%s) didn't match expected output: %s", 278 | serialized, expected) 279 | } 280 | } 281 | 282 | func TestTagSort(t *testing.T) { 283 | contains := func(key string, tags TagSet) bool { 284 | for _, t := range tags { 285 | if t.Key == key { 286 | return true 287 | } 288 | } 289 | return false 290 | } 291 | 292 | for n := 0; n < 20; n++ { 293 | tags := randomTagSet(t, "v", n) 294 | keys := make([]string, 0, len(tags)+5) 295 | for _, t := range tags { 296 | keys = append(keys, t.Key) 297 | } 298 | for i := 0; i < 5; i++ { 299 | for { 300 | s := randomString(t, 10) 301 | if !contains(s, tags) { 302 | keys = append(keys, randomString(t, 10)) 303 | break 304 | } 305 | } 306 | } 307 | for _, key := range keys { 308 | i := tags.Search(key) 309 | j := sort.Search(len(tags), func(i int) bool { 310 | return tags[i].Key >= key 311 | }) 312 | if i != j { 313 | t.Errorf("%d: Search got: %d want: %d", n, i, j) 314 | } 315 | } 316 | 317 | for _, key := range keys { 318 | exp := contains(key, tags) 319 | got := tags.Contains(key) 320 | if exp != got { 321 | t.Errorf("%d: tags contains (%q) want: %t got: %t", n, key, exp, got) 322 | } 323 | } 324 | 325 | for i := range tags { 326 | j := tags.Search(tags[i].Key) 327 | if j != i { 328 | t.Errorf("%d: search did not find %q-%d: %d", n, tags[i].Key, i, j) 329 | } 330 | } 331 | } 332 | } 333 | 334 | func randomTagSet(t testing.TB, valPrefix string, size int) TagSet { 335 | s := make(TagSet, size) 336 | for i := 0; i < len(s); i++ { 337 | s[i] = Tag{ 338 | Key: randomString(t, 32), 339 | Value: fmt.Sprintf("%s%d", valPrefix, i), 340 | } 341 | } 342 | s.Sort() 343 | return s 344 | } 345 | 346 | func TestTagInsert(t *testing.T) { 347 | t1 := randomTagSet(t, "t1_", 1000) 348 | t2 := randomTagSet(t, "t2_", 1000) 349 | if !sort.IsSorted(t1) { 350 | t.Fatal("tags being inserted into must be sorted!") 351 | } 352 | for i := range t2 { 353 | t1 = t1.Insert(t2[i]) 354 | if !sort.IsSorted(t1) { 355 | t.Fatalf("%d: inserting tag failed: %+v", i, t2[i]) 356 | } 357 | } 358 | 359 | // Make sure Insert is COW 360 | 361 | // If the tag we're inserting is already in the set we 362 | // should *not* return a copy 363 | t.Run("NoCopy", func(t *testing.T) { 364 | t1 := TagSet{{"k", "v"}} 365 | t2 := t1.Insert(Tag{"k", "v"}) 366 | if &t1[0] != &t2[0] { 367 | t.Errorf("Copy: %p -- %p", 368 | (*uintptr)(unsafe.Pointer(&t1[0])), 369 | (*uintptr)(unsafe.Pointer(&t2[0])), 370 | ) 371 | } 372 | }) 373 | 374 | t.Run("CopyEqualKey", func(t *testing.T) { 375 | t1 := TagSet{{"k", "v1"}} 376 | t2 := t1.Insert(Tag{"k", "v2"}) 377 | if &t1[0] == &t2[0] { 378 | t.Errorf("Copy: %p -- %p", 379 | (*uintptr)(unsafe.Pointer(&t1[0])), 380 | (*uintptr)(unsafe.Pointer(&t2[0])), 381 | ) 382 | } 383 | }) 384 | 385 | t.Run("Copy", func(t *testing.T) { 386 | t1 := make(TagSet, 0, 2) 387 | t1 = append(t1, Tag{"k1", "v1"}) 388 | t2 := t1.Insert(Tag{"k2", "v2"}) 389 | if &t1[0] == &t2[0] { 390 | t.Errorf("TagSet was modified, but not copied: %p -- %p", 391 | (*uintptr)(unsafe.Pointer(&t1[0])), 392 | (*uintptr)(unsafe.Pointer(&t2[0])), 393 | ) 394 | } 395 | }) 396 | } 397 | 398 | func mergeTagSetsReference(s1, s2 TagSet) TagSet { 399 | seen := make(map[string]bool) 400 | var a TagSet 401 | for _, t := range s2 { 402 | a = append(a, t) 403 | seen[t.Key] = true 404 | } 405 | for _, t := range s1 { 406 | if !seen[t.Key] { 407 | a = append(a, t) 408 | } 409 | } 410 | a.Sort() 411 | return a 412 | } 413 | 414 | func makeScratch(s1, s2 TagSet) (TagSet, TagSet) { 415 | a := make(TagSet, len(s1)+len(s2)) 416 | copy(a[len(s1):], s2) 417 | return a, a[len(s1):] 418 | } 419 | 420 | func tagSetEqual(s1, s2 TagSet) bool { 421 | if len(s1) != len(s2) { 422 | return false 423 | } 424 | for i := range s1 { 425 | if s1[i] != s2[i] { 426 | return false 427 | } 428 | } 429 | return true 430 | } 431 | 432 | func TestMergeTagSets(t *testing.T) { 433 | for i := 0; i < 100; i++ { 434 | s1 := randomTagSet(t, "s1_", i) 435 | s2 := randomTagSet(t, "s2_", i) 436 | for i := 0; i < len(s2); i++ { 437 | if i&1 == 0 { 438 | s2[i] = s1[i] 439 | } 440 | } 441 | s1.Sort() 442 | s2.Sort() 443 | if !sort.IsSorted(s1) { 444 | t.Fatal("s1 not sorted") 445 | } 446 | if !sort.IsSorted(s2) { 447 | t.Fatal("s2 not sorted") 448 | } 449 | 450 | expected := mergeTagSetsReference(s1, s2) 451 | a, to := makeScratch(s1, s2) 452 | got := mergeTagSets(s1, to, a) 453 | if !sort.IsSorted(got) { 454 | t.Errorf("merging %d tagSets failed: not sorted", i) 455 | } 456 | if !tagSetEqual(got, expected) { 457 | // t.Errorf("merging %d tagSets", i) 458 | t.Errorf("merging %d tagSets failed\n# Got:\n%+v\n# Want:\n%+v\n# S1:\n%+v\n# S2:\n%+v\n", 459 | i, got, expected, s1, s2) 460 | } 461 | } 462 | } 463 | 464 | func TestSerializeTagSet(t *testing.T) { 465 | if s := TagSet(nil).Serialize("name"); s != "name" { 466 | t.Errorf("got: %q want: %q", s, "name") 467 | } 468 | 469 | const exp = "name.__k1=v1.__k2=v2" 470 | set := NewTagSet(map[string]string{ 471 | "k1": "v1", 472 | "k2": "v2", 473 | }) 474 | if s := set.Serialize("name"); s != exp { 475 | t.Errorf("got: %q want: %q", s, "name") 476 | } 477 | } 478 | 479 | func TestNewTag(t *testing.T) { 480 | exp := Tag{ 481 | Key: "key", 482 | Value: ReplaceChars("value.a:b|c"), 483 | } 484 | got := NewTag("key", "value.a:b|c") 485 | if got != exp { 486 | t.Errorf("NewTag: got: %+v want: %+v", got, exp) 487 | } 488 | } 489 | 490 | func TestNewTagSet(t *testing.T) { 491 | m := map[string]string{ 492 | "x": "x", 493 | "y": "y", 494 | "z": "z", 495 | "c": "c|", 496 | "b": "b:", 497 | "a": "a.", 498 | "": "empty", 499 | "empty": "", 500 | } 501 | exp := TagSet{ 502 | {"a", "a_"}, 503 | {"b", "b_"}, 504 | {"c", "c_"}, 505 | {"x", "x"}, 506 | {"y", "y"}, 507 | {"z", "z"}, 508 | } 509 | got := NewTagSet(m) 510 | if !reflect.DeepEqual(got, exp) { 511 | t.Errorf("NewTagSet: got: %+v want: %+v", got, exp) 512 | } 513 | } 514 | 515 | func benchmarkSerializeTags(b *testing.B, n int) { 516 | const name = "prefix" 517 | tags := make(map[string]string, n) 518 | for i := 0; i < n; i++ { 519 | k := fmt.Sprintf("key%d", i) 520 | v := fmt.Sprintf("val%d", i) 521 | tags[k] = v 522 | } 523 | b.ResetTimer() 524 | for i := 0; i < b.N; i++ { 525 | SerializeTags(name, tags) 526 | } 527 | } 528 | 529 | func BenchmarkSerializeTags(b *testing.B) { 530 | for i := 1; i <= 10; i++ { 531 | b.Run(fmt.Sprintf("%d", i), func(b *testing.B) { 532 | benchmarkSerializeTags(b, i) 533 | }) 534 | } 535 | } 536 | 537 | func benchmarkSerializeTagSet(b *testing.B, n int) { 538 | const name = "prefix" 539 | tags := make(TagSet, 0, n) 540 | for i := 0; i < n; i++ { 541 | k := fmt.Sprintf("key%d", i) 542 | v := fmt.Sprintf("val%d", i) 543 | tags = append(tags, Tag{ 544 | Key: k, 545 | Value: v, 546 | }) 547 | } 548 | tags.Sort() 549 | b.ResetTimer() 550 | for i := 0; i < b.N; i++ { 551 | tags.Serialize(name) 552 | } 553 | } 554 | 555 | func BenchmarkSerializeTagSet(b *testing.B) { 556 | for i := 1; i <= 10; i++ { 557 | b.Run(fmt.Sprintf("%d", i), func(b *testing.B) { 558 | benchmarkSerializeTagSet(b, i) 559 | }) 560 | } 561 | } 562 | 563 | // TODO (CEV): consider removing this 564 | func BenchmarkTagSearch(b *testing.B) { 565 | rr := rand.New(rand.NewSource(12345)) 566 | tags := make(TagSet, 5) 567 | for i := range tags { 568 | tags[i].Key = strconv.FormatInt(rr.Int63(), 10) 569 | tags[i].Value = strconv.FormatInt(rr.Int63(), 10) 570 | } 571 | tags.Sort() 572 | 573 | for i := 0; i < b.N; i++ { 574 | for _, tag := range tags { 575 | _ = tags.Search(tag.Key) 576 | } 577 | } 578 | } 579 | 580 | // TODO (CEV): consider removing this 581 | func BenchmarkTagSearch_Reference(b *testing.B) { 582 | rr := rand.New(rand.NewSource(12345)) 583 | tags := make(TagSet, 5) 584 | for i := range tags { 585 | tags[i].Key = strconv.FormatInt(rr.Int63(), 10) 586 | tags[i].Value = strconv.FormatInt(rr.Int63(), 10) 587 | } 588 | tags.Sort() 589 | 590 | for i := 0; i < b.N; i++ { 591 | for _, tag := range tags { 592 | _ = sort.Search(len(tags), func(i int) bool { 593 | return tags[i].Key >= tag.Key 594 | }) 595 | } 596 | } 597 | } 598 | 599 | func benchTagSort(b *testing.B, size int) { 600 | // CEV: this isn't super accurate since we also time 601 | // the copying the orig slice into the test slice, 602 | // but its still useful. 603 | 604 | // use a fixed source so that results are comparable 605 | rr := rand.New(rand.NewSource(12345)) 606 | orig := make(TagSet, size) 607 | for i := range orig { 608 | orig[i].Key = strconv.FormatInt(rr.Int63(), 10) 609 | orig[i].Value = strconv.FormatInt(rr.Int63(), 10) 610 | } 611 | rr.Shuffle(len(orig), func(i, j int) { 612 | orig[i], orig[j] = orig[j], orig[i] 613 | }) 614 | tags := make(TagSet, len(orig)) 615 | 616 | b.ResetTimer() 617 | for i := 0; i < b.N; i++ { 618 | copy(tags, orig) 619 | tags.Sort() 620 | } 621 | } 622 | 623 | func BenchmarkTagSort(b *testing.B) { 624 | if testing.Short() { 625 | b.Skip("short test") 626 | } 627 | for i := 2; i <= 10; i++ { 628 | b.Run(fmt.Sprint(i), func(b *testing.B) { 629 | benchTagSort(b, i) 630 | }) 631 | } 632 | } 633 | 634 | // TODO: consider making this benchmark smaller 635 | func BenchmarkMergeTagSets(b *testing.B) { 636 | if testing.Short() { 637 | b.Skip("short test") 638 | } 639 | t1 := make(TagSet, 10) 640 | t2 := make(TagSet, 10) 641 | for i := 0; i < len(t1); i++ { 642 | t1[i] = Tag{ 643 | Key: fmt.Sprintf("k1%d", i), 644 | Value: fmt.Sprintf("v1_%d", i), 645 | } 646 | t2[i] = Tag{ 647 | Key: fmt.Sprintf("k2%d", i), 648 | Value: fmt.Sprintf("v2_%d", i), 649 | } 650 | } 651 | t1.Sort() 652 | t2.Sort() 653 | 654 | scratch := make(TagSet, len(t1)+len(t2)) 655 | 656 | b.ResetTimer() 657 | b.Run("KeysNotEqual", func(b *testing.B) { 658 | for size := 2; size <= 10; size += 2 { 659 | b.Run(fmt.Sprint(size), func(b *testing.B) { 660 | s1 := t1[:size] 661 | s2 := t2[:size] 662 | for i := 0; i < b.N; i++ { 663 | mergeTagSets(s1, s2, scratch) 664 | } 665 | }) 666 | } 667 | }) 668 | 669 | b.Run("KeysHalfEqual", func(b *testing.B) { 670 | for i := range t2 { 671 | if i&1 != 0 { 672 | t2[i].Key = t1[i].Key 673 | } 674 | } 675 | t2.Sort() 676 | for size := 2; size <= 10; size += 2 { 677 | b.Run(fmt.Sprint(size), func(b *testing.B) { 678 | s1 := t1[:size] 679 | s2 := t2[:size] 680 | for i := 0; i < b.N; i++ { 681 | mergeTagSets(s1, s2, scratch) 682 | } 683 | }) 684 | } 685 | }) 686 | 687 | b.Run("KeysEqual", func(b *testing.B) { 688 | for i := range t2 { 689 | t2[i].Key = t1[i].Key 690 | } 691 | for size := 2; size <= 10; size += 2 { 692 | b.Run(fmt.Sprint(size), func(b *testing.B) { 693 | s1 := t1[:size] 694 | s2 := t2[:size] 695 | for i := 0; i < b.N; i++ { 696 | mergeTagSets(s1, s2, scratch) 697 | } 698 | }) 699 | } 700 | }) 701 | } 702 | 703 | func BenchmarkTagSetSearch_Reference(b *testing.B) { 704 | var keys [5]string 705 | tags := randomTagSet(b, "v_", len(keys)) 706 | for i := 0; i < len(tags); i++ { 707 | keys[i] = tags[i].Key 708 | } 709 | b.ResetTimer() 710 | for i := 0; i < b.N; i++ { 711 | key := keys[i%len(keys)] 712 | sort.Search(len(tags), func(i int) bool { 713 | return tags[i].Key >= key 714 | }) 715 | } 716 | } 717 | 718 | func BenchmarkTagSetSearch(b *testing.B) { 719 | var keys [5]string 720 | tags := randomTagSet(b, "v_", len(keys)) 721 | for i := 0; i < len(tags); i++ { 722 | keys[i] = tags[i].Key 723 | } 724 | b.ResetTimer() 725 | for i := 0; i < b.N; i++ { 726 | tags.Search(keys[i%len(keys)]) 727 | } 728 | } 729 | 730 | /////////////////////////////////////////////////////////////////// 731 | // TagSet Tests 732 | 733 | func mergeTagsReference(set TagSet, tags map[string]string) TagSet { 734 | a := make(TagSet, 0, len(tags)) 735 | for k, v := range tags { 736 | if k != "" && v != "" { 737 | a = append(a, Tag{Key: k, Value: ReplaceChars(v)}) 738 | } 739 | } 740 | return mergeTagSetsReference(set, a) 741 | } 742 | 743 | func tagMapsEqual(m1, m2 map[string]string) bool { 744 | if len(m1) != len(m2) { 745 | return false 746 | } 747 | for k, v := range m1 { 748 | if vv, ok := m2[k]; !ok || vv != v { 749 | return false 750 | } 751 | } 752 | return true 753 | } 754 | 755 | func testMergeTags(t *testing.T, s1, s2 TagSet, perInstanceTag bool) { 756 | tags := make(map[string]string, len(s2)) 757 | origTags := make(map[string]string, len(s2)) 758 | for _, p := range s2 { 759 | tags[p.Key] = p.Value 760 | origTags[p.Key] = p.Value 761 | } 762 | 763 | set := s1 764 | origPairs := append(TagSet(nil), set...) 765 | expected := mergeTagsReference(set, tags) 766 | 767 | var got TagSet 768 | if perInstanceTag { 769 | if !expected.Contains("_f") { 770 | expected = expected.Insert(Tag{Key: "_f", Value: "i"}) 771 | } 772 | got = set.MergePerInstanceTags(tags) 773 | } else { 774 | got = set.MergeTags(tags) 775 | } 776 | if !tagSetEqual(got, expected) { 777 | t.Errorf("{s1=%d, s2=%d}: bad merge:\n# Got:\n%+v\n\n# Want:\n%+v\n", 778 | len(s1), len(s2), got, expected) 779 | } 780 | if !tagSetEqual(set, origPairs) { 781 | t.Fatalf("scope tags modified:\n# Got:\n%+v\n\n# Want:\n%+v\n", 782 | set, origPairs) 783 | } 784 | if !tagMapsEqual(tags, origTags) { 785 | t.Fatalf("tag map modified:\n# Got:\n%v\n\n# Want:\n%v\n", 786 | set, origPairs) 787 | } 788 | if perInstanceTag { 789 | if !got.Contains("_f") { 790 | t.Fatal("missing per-instance tag") 791 | } 792 | exp := "i" // default 793 | for _, p := range expected { 794 | if p.Key == "_f" { 795 | exp = ReplaceChars(p.Value) 796 | break 797 | } 798 | } 799 | tag := got[got.Search("_f")].Value 800 | if tag != exp { 801 | t.Fatalf("per-instance tag want: %q got: %q: %+v", exp, tag, got) 802 | } 803 | } 804 | } 805 | 806 | func TestMergePerInstanceTags(t *testing.T) { 807 | t.Parallel() 808 | 809 | rr := rand.New(rand.NewSource(time.Now().UnixNano())) 810 | 811 | type testCase struct { 812 | n1, n2 int 813 | hasPerInstanceTag bool 814 | } 815 | tests := make([]testCase, 0, 2100) 816 | 817 | // make sure we cover all 0..2 test cases 818 | for i := 0; i <= 2; i++ { 819 | for j := 0; j <= 2; j++ { 820 | for k := 0; k < 2; k++ { 821 | tests = append(tests, testCase{i, j, k == 1}) 822 | } 823 | } 824 | } 825 | // add a whole bunch of random cases 826 | for i := 0; i < 2000; i++ { 827 | tests = append(tests, testCase{rr.Intn(8), rr.Intn(8), false}) 828 | } 829 | 830 | for _, x := range tests { 831 | s1 := randomTagSet(t, "v", x.n1) 832 | s2 := randomTagSet(t, "v", x.n2) 833 | if x.hasPerInstanceTag { 834 | s1 = s1.Insert(Tag{Key: "_f", Value: "foo"}) 835 | } 836 | if rr.Float64() < 0.1 { 837 | s1 = s1.Insert(Tag{Key: "_f", Value: "foo"}) 838 | } 839 | if rr.Float64() < 0.1 { 840 | s2 = s2.Insert(Tag{Key: "_f", Value: "bar"}) 841 | } 842 | 843 | // Add some invalid chars to s2 844 | for i := range s2 { 845 | if rr.Float64() < 0.2 { 846 | s2[i].Value += "|" 847 | } 848 | if rr.Float64() < 0.1 { 849 | s2[i].Value = "" 850 | } 851 | if rr.Float64() < 0.1 { 852 | s2[i].Key = "" 853 | } 854 | } 855 | testMergeTags(t, s1, s2, true) 856 | } 857 | } 858 | 859 | func TestMergeTags(t *testing.T) { 860 | t.Parallel() 861 | 862 | rr := rand.New(rand.NewSource(time.Now().UnixNano())) 863 | 864 | type testCase struct { 865 | n1, n2 int 866 | } 867 | tests := make([]testCase, 0, 2100) 868 | 869 | // make sure we cover all 0..2 test cases 870 | for i := 0; i <= 2; i++ { 871 | for j := 0; j <= 2; j++ { 872 | tests = append(tests, testCase{i, j}) 873 | } 874 | } 875 | // add a whole bunch of random cases 876 | for i := 0; i < 2000; i++ { 877 | tests = append(tests, testCase{rr.Intn(64), rr.Intn(64)}) 878 | } 879 | 880 | for _, x := range tests { 881 | s1 := randomTagSet(t, "v", x.n1) 882 | s2 := randomTagSet(t, "v", x.n2) 883 | 884 | // Add some invalid chars to s2 885 | for i := range s2 { 886 | if rr.Float64() < 0.2 { 887 | s2[i].Value += "|" 888 | } 889 | if rr.Float64() < 0.1 { 890 | s2[i].Value = "" 891 | } 892 | if rr.Float64() < 0.1 { 893 | s2[i].Key = "" 894 | } 895 | } 896 | testMergeTags(t, s1, s2, false) 897 | } 898 | } 899 | 900 | func TestMergeOneTagPanic(t *testing.T) { 901 | tags := map[string]string{ 902 | "k1": "v1", 903 | "k2": "v2", 904 | } 905 | defer func() { 906 | if recover() == nil { 907 | t.Fatal("expected panic got none") 908 | } 909 | }() 910 | mergeOneTag(TagSet{}, tags) 911 | } 912 | 913 | func randomString(tb testing.TB, size int) string { 914 | b := make([]byte, hex.DecodedLen(size)) 915 | if _, err := crand.Read(b); err != nil { 916 | tb.Fatal(err) 917 | } 918 | return hex.EncodeToString(b) 919 | } 920 | 921 | // TODO: rename this once we rename the mergePairs method 922 | func benchScopeMergeTags(b *testing.B, baseSize, tagsSize int) { 923 | set := make(TagSet, baseSize) 924 | for i := range set { 925 | set[i] = Tag{ 926 | Key: fmt.Sprintf("key1_%d", i), 927 | Value: fmt.Sprintf("val1_%d", i), 928 | } 929 | } 930 | tags := make(map[string]string, tagsSize) 931 | for i := 0; i < tagsSize; i++ { 932 | key := fmt.Sprintf("key2_%d", i) 933 | val := fmt.Sprintf("val2_%d", i) 934 | tags[key] = val 935 | } 936 | b.ResetTimer() 937 | 938 | for i := 0; i < b.N; i++ { 939 | set.MergeTags(tags) 940 | } 941 | } 942 | 943 | // TODO: rename this once we rename the mergePairs method 944 | func BenchmarkScopeMergeTags(b *testing.B) { 945 | if testing.Short() { 946 | b.Skip("short test") 947 | } 948 | for baseSize := 1; baseSize <= 8; baseSize++ { 949 | for tagSize := 1; tagSize <= 8; tagSize++ { 950 | b.Run(fmt.Sprintf("%d_%d", baseSize, tagSize), func(b *testing.B) { 951 | benchScopeMergeTags(b, baseSize, tagSize) 952 | }) 953 | } 954 | } 955 | } 956 | 957 | /////////////////////////////////////////////////////////////////// 958 | // Parse Tests 959 | 960 | var parseTagsTests = []struct { 961 | Stat string 962 | Name string 963 | Tags map[string]string 964 | }{ 965 | { 966 | Stat: "", 967 | Name: "", 968 | }, 969 | { 970 | Stat: "prefix.name", 971 | Name: "prefix.name", 972 | }, 973 | { 974 | Stat: "prefix.c.___f=i.__tag1=value1:12|c\n", 975 | Name: "prefix.c", 976 | Tags: map[string]string{ 977 | "_f": "i", 978 | "tag1": "value1", 979 | }, 980 | }, 981 | // malformed 982 | { 983 | Stat: "invalid.__tag1=value1.__tag2", 984 | Name: "invalid", 985 | Tags: map[string]string{ 986 | "tag1": "value1", 987 | }, 988 | }, 989 | { 990 | Stat: "invalid.__tag1.__tag2=value2", 991 | Name: "invalid", 992 | Tags: map[string]string{ 993 | "tag2": "value2", 994 | }, 995 | }, 996 | { 997 | Stat: "invalid.__tag1=.__tag2=value2", 998 | Name: "invalid", 999 | Tags: map[string]string{ 1000 | "tag1": "", 1001 | "tag2": "value2", 1002 | }, 1003 | }, 1004 | } 1005 | 1006 | func TestRemoveStatValue(t *testing.T) { 1007 | tests := map[string]string{ 1008 | "": "", 1009 | "a": "a", 1010 | "a:b:12|c\n": "a:b", // likely invalid, but handle it 1011 | "prefix.c.___f=i.__tag1=value1:12|c\n": "prefix.c.___f=i.__tag1=value1", 1012 | "prefix.c.___f=i.__tag1=value1": "prefix.c.___f=i.__tag1=value1", 1013 | } 1014 | for in, exp := range tests { 1015 | got := removeStatValue(in) 1016 | if got != exp { 1017 | t.Errorf("%q: got: %q want: %q", in, got, exp) 1018 | } 1019 | } 1020 | } 1021 | 1022 | func TestParseTags(t *testing.T) { 1023 | for _, x := range parseTagsTests { 1024 | s, m := ParseTags(x.Stat) 1025 | if s != x.Name { 1026 | t.Errorf("%+v: Name: got: %q want: %q", x, s, x.Name) 1027 | } 1028 | if !tagMapsEqual(m, x.Tags) { 1029 | t.Errorf("%+v: Tags: got: %q want: %q", x, m, x.Tags) 1030 | } 1031 | } 1032 | } 1033 | 1034 | func TestParseTagSet(t *testing.T) { 1035 | for _, x := range parseTagsTests { 1036 | s, set := ParseTagSet(x.Stat) 1037 | if s != x.Name { 1038 | t.Errorf("%+v: Name: got: %q want: %q", x, s, x.Name) 1039 | } 1040 | exp := make(TagSet, 0, len(x.Tags)) 1041 | for k, v := range x.Tags { 1042 | exp = append(exp, Tag{Key: k, Value: v}) 1043 | } 1044 | exp.Sort() 1045 | if !tagSetEqual(set, exp) { 1046 | t.Errorf("%+v: Tags: got: %q want: %q", x, set, exp) 1047 | } 1048 | } 1049 | } 1050 | --------------------------------------------------------------------------------