github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/server/debug/pprofui/storage_mem.go (about) 1 // Copyright 2018 The Cockroach Authors. 2 // 3 // Use of this software is governed by the Business Source License 4 // included in the file licenses/BSL.txt. 5 // 6 // As of the Change Date specified in that file, in accordance with 7 // the Business Source License, use of this software will be governed 8 // by the Apache License, Version 2.0, included in the file 9 // licenses/APL.txt. 10 11 package pprofui 12 13 import ( 14 "bytes" 15 "fmt" 16 "io" 17 "sort" 18 "sync/atomic" 19 "time" 20 21 "github.com/cockroachdb/cockroach/pkg/util/syncutil" 22 "github.com/cockroachdb/cockroach/pkg/util/timeutil" 23 "github.com/cockroachdb/errors" 24 ) 25 26 type record struct { 27 id string 28 t time.Time 29 b []byte 30 } 31 32 // A MemStorage is a Storage implementation that holds recent profiles in memory. 33 type MemStorage struct { 34 mu struct { 35 syncutil.Mutex 36 records []record // sorted by record.t 37 } 38 idGen int32 // accessed atomically 39 keepDuration time.Duration // zero for disabled 40 keepNumber int // zero for disabled 41 } 42 43 var _ Storage = &MemStorage{} 44 45 // NewMemStorage creates a MemStorage that retains the most recent n records 46 // as long as they are less than d old. 47 // 48 // Records are dropped only when there is activity (i.e. an old record will 49 // only be dropped the next time the storage is accessed). 50 func NewMemStorage(n int, d time.Duration) *MemStorage { 51 return &MemStorage{ 52 keepNumber: n, 53 keepDuration: d, 54 } 55 } 56 57 // ID implements Storage. 58 func (s *MemStorage) ID() string { 59 return fmt.Sprint(atomic.AddInt32(&s.idGen, 1)) 60 } 61 62 func (s *MemStorage) cleanLocked() { 63 if l, m := len(s.mu.records), s.keepNumber; l > m && m != 0 { 64 s.mu.records = append([]record(nil), s.mu.records[l-m:]...) 65 } 66 now := timeutil.Now() 67 if pos := sort.Search(len(s.mu.records), func(i int) bool { 68 return s.mu.records[i].t.Add(s.keepDuration).After(now) 69 }); pos < len(s.mu.records) && s.keepDuration != 0 { 70 s.mu.records = append([]record(nil), s.mu.records[pos:]...) 71 } 72 } 73 74 // Store implements Storage. 75 func (s *MemStorage) Store(id string, write func(io.Writer) error) error { 76 var b bytes.Buffer 77 if err := write(&b); err != nil { 78 return err 79 } 80 s.mu.Lock() 81 defer s.mu.Unlock() 82 s.mu.records = append(s.mu.records, record{id: id, t: timeutil.Now(), b: b.Bytes()}) 83 sort.Slice(s.mu.records, func(i, j int) bool { 84 return s.mu.records[i].t.Before(s.mu.records[j].t) 85 }) 86 s.cleanLocked() 87 return nil 88 } 89 90 // Get implements Storage. 91 func (s *MemStorage) Get(id string, read func(io.Reader) error) error { 92 s.mu.Lock() 93 defer s.mu.Unlock() 94 for _, v := range s.mu.records { 95 if v.id == id { 96 return read(bytes.NewReader(v.b)) 97 } 98 } 99 return errors.Errorf("profile not found; it may have expired") 100 }