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  }