| // Copyright 2016 Google Inc. All Rights Reserved. |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| package trace |
| |
| import ( |
| crand "crypto/rand" |
| "encoding/binary" |
| "fmt" |
| "math/rand" |
| "sync" |
| "time" |
| |
| "golang.org/x/time/rate" |
| ) |
| |
| type SamplingPolicy interface { |
| // Sample determines whether to sample the next request. If so, it also |
| // returns a string and rate describing the reason the request was chosen. |
| Sample() (sample bool, policy string, rate float64) |
| } |
| |
| type sampler struct { |
| fraction float64 |
| *rate.Limiter |
| *rand.Rand |
| sync.Mutex |
| } |
| |
| func (s *sampler) Sample() (sample bool, reason string, rate float64) { |
| s.Lock() |
| x := s.Float64() |
| s.Unlock() |
| return s.sample(time.Now(), x) |
| } |
| |
| // sample contains the a deterministic, time-independent logic of Sample. |
| func (s *sampler) sample(now time.Time, x float64) (bool, string, float64) { |
| if x >= s.fraction || !s.AllowN(now, 1) { |
| return false, "", 0.0 |
| } |
| if s.fraction < 1.0 { |
| return true, "fraction", s.fraction |
| } |
| return true, "qps", float64(s.Limit()) |
| } |
| |
| // NewLimitedSampler returns a sampling policy that traces a given fraction of |
| // requests, and enforces a limit on the number of traces per second. |
| // Returns a nil SamplingPolicy if either fraction or maxqps is zero. |
| func NewLimitedSampler(fraction, maxqps float64) (SamplingPolicy, error) { |
| if !(fraction >= 0) { |
| return nil, fmt.Errorf("invalid fraction %f", fraction) |
| } |
| if !(maxqps >= 0) { |
| return nil, fmt.Errorf("invalid maxqps %f", maxqps) |
| } |
| if fraction == 0 || maxqps == 0 { |
| return nil, nil |
| } |
| // Set a limit on the number of accumulated "tokens", to limit bursts of |
| // traced requests. Use one more than a second's worth of tokens, or 100, |
| // whichever is smaller. |
| // See https://godoc.org/golang.org/x/time/rate#NewLimiter. |
| maxTokens := 100 |
| if maxqps < 99.0 { |
| maxTokens = 1 + int(maxqps) |
| } |
| var seed int64 |
| if err := binary.Read(crand.Reader, binary.LittleEndian, &seed); err != nil { |
| seed = time.Now().UnixNano() |
| } |
| s := sampler{ |
| fraction: fraction, |
| Limiter: rate.NewLimiter(rate.Limit(maxqps), maxTokens), |
| Rand: rand.New(rand.NewSource(seed)), |
| } |
| return &s, nil |
| } |