1// Copyright 2016 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package stat
16
17import (
18	"bytes"
19	"encoding/csv"
20	"fmt"
21	"io"
22	"math"
23	"sort"
24	"strconv"
25	"text/tabwriter"
26	"time"
27)
28
29type byDuration []time.Duration
30
31func (data byDuration) Len() int           { return len(data) }
32func (data byDuration) Swap(i, j int)      { data[i], data[j] = data[j], data[i] }
33func (data byDuration) Less(i, j int) bool { return data[i] < data[j] }
34
35// quantile returns a value representing the kth of q quantiles.
36// May alter the order of data.
37func quantile(data []time.Duration, k, q int) (quantile time.Duration, ok bool) {
38	if len(data) < 1 {
39		return 0, false
40	}
41	if k > q {
42		return 0, false
43	}
44	if k < 0 || q < 1 {
45		return 0, false
46	}
47
48	sort.Sort(byDuration(data))
49
50	if k == 0 {
51		return data[0], true
52	}
53	if k == q {
54		return data[len(data)-1], true
55	}
56
57	bucketSize := float64(len(data)-1) / float64(q)
58	i := float64(k) * bucketSize
59
60	lower := int(math.Trunc(i))
61	var upper int
62	if i > float64(lower) && lower+1 < len(data) {
63		// If the quantile lies between two elements
64		upper = lower + 1
65	} else {
66		upper = lower
67	}
68	weightUpper := i - float64(lower)
69	weightLower := 1 - weightUpper
70	return time.Duration(weightLower*float64(data[lower]) + weightUpper*float64(data[upper])), true
71}
72
73// Aggregate is an aggregate of latencies.
74type Aggregate struct {
75	Name               string
76	Count, Errors      int
77	Min, Median, Max   time.Duration
78	P75, P90, P95, P99 time.Duration // percentiles
79}
80
81// NewAggregate constructs an aggregate from latencies. Returns nil if latencies does not contain aggregateable data.
82func NewAggregate(name string, latencies []time.Duration, errorCount int) *Aggregate {
83	agg := Aggregate{Name: name, Count: len(latencies), Errors: errorCount}
84
85	if len(latencies) == 0 {
86		return nil
87	}
88	var ok bool
89	if agg.Min, ok = quantile(latencies, 0, 2); !ok {
90		return nil
91	}
92	if agg.Median, ok = quantile(latencies, 1, 2); !ok {
93		return nil
94	}
95	if agg.Max, ok = quantile(latencies, 2, 2); !ok {
96		return nil
97	}
98	if agg.P75, ok = quantile(latencies, 75, 100); !ok {
99		return nil
100	}
101	if agg.P90, ok = quantile(latencies, 90, 100); !ok {
102		return nil
103	}
104	if agg.P95, ok = quantile(latencies, 95, 100); !ok {
105		return nil
106	}
107	if agg.P99, ok = quantile(latencies, 99, 100); !ok {
108		return nil
109	}
110	return &agg
111}
112
113func (agg *Aggregate) String() string {
114	if agg == nil {
115		return "no data"
116	}
117	var buf bytes.Buffer
118	tw := tabwriter.NewWriter(&buf, 0, 0, 1, ' ', 0) // one-space padding
119	fmt.Fprintf(tw, "min:\t%v\nmedian:\t%v\nmax:\t%v\n95th percentile:\t%v\n99th percentile:\t%v\n",
120		agg.Min, agg.Median, agg.Max, agg.P95, agg.P99)
121	tw.Flush()
122	return buf.String()
123}
124
125// WriteCSV writes a csv file to the given Writer,
126// with a header row and one row per aggregate.
127func WriteCSV(aggs []*Aggregate, iow io.Writer) (err error) {
128	w := csv.NewWriter(iow)
129	defer func() {
130		w.Flush()
131		if err == nil {
132			err = w.Error()
133		}
134	}()
135	err = w.Write([]string{"name", "count", "errors", "min", "median", "max", "p75", "p90", "p95", "p99"})
136	if err != nil {
137		return err
138	}
139	for _, agg := range aggs {
140		err = w.Write([]string{
141			agg.Name, strconv.Itoa(agg.Count), strconv.Itoa(agg.Errors),
142			agg.Min.String(), agg.Median.String(), agg.Max.String(),
143			agg.P75.String(), agg.P90.String(), agg.P95.String(), agg.P99.String(),
144		})
145		if err != nil {
146			return err
147		}
148	}
149	return nil
150}
151