github.com/grafana/pyroscope@v1.18.0/pkg/metastore/fsm/snapshot.go (about)

     1  package fsm
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"context"
     7  	"io"
     8  	"runtime/pprof"
     9  	"time"
    10  
    11  	"github.com/go-kit/log"
    12  	"github.com/go-kit/log/level"
    13  	"github.com/hashicorp/raft"
    14  	"github.com/klauspost/compress/zstd"
    15  	"go.etcd.io/bbolt"
    16  
    17  	"github.com/grafana/pyroscope/pkg/util/ratelimit"
    18  )
    19  
    20  type snapshotWriter struct {
    21  	logger      log.Logger
    22  	tx          *bbolt.Tx
    23  	metrics     *metrics
    24  	compression string
    25  	rate        int
    26  }
    27  
    28  func (s *snapshotWriter) Persist(sink raft.SnapshotSink) (err error) {
    29  	ctx := context.Background()
    30  	pprof.SetGoroutineLabels(pprof.WithLabels(ctx, pprof.Labels("metastore_op", "persist")))
    31  	defer pprof.SetGoroutineLabels(ctx)
    32  
    33  	start := time.Now()
    34  	level.Debug(s.logger).Log("msg", "persisting snapshot", "sink_id", sink.ID())
    35  	defer func() {
    36  		s.metrics.boltDBPersistSnapshotDuration.Observe(time.Since(start).Seconds())
    37  		if err == nil {
    38  			level.Info(s.logger).Log("msg", "persisted snapshot", "sink_id", sink.ID(), "duration", time.Since(start))
    39  			if err = sink.Close(); err != nil {
    40  				level.Error(s.logger).Log("msg", "failed to close sink", "err", err)
    41  			}
    42  			return
    43  		}
    44  		level.Error(s.logger).Log("msg", "failed to persist snapshot", "err", err)
    45  		if err = sink.Cancel(); err != nil {
    46  			level.Error(s.logger).Log("msg", "failed to cancel snapshot sink", "err", err)
    47  		}
    48  	}()
    49  
    50  	// Needed to measure the actual snapshot size written to the sink.
    51  	w := &writeCounter{Writer: sink}
    52  	var dst io.Writer = w
    53  
    54  	// Optionally rate limit the snapshot writes to avoid overwhelming the disk.
    55  	// This is useful when the snapshot is written to the same disk as the WAL.
    56  	if s.rate > 0 {
    57  		const minRate = 10 // 10 MB/s
    58  		if s.rate < minRate {
    59  			s.rate = minRate
    60  		}
    61  		s.rate <<= 20
    62  		dst = ratelimit.NewWriter(dst, ratelimit.NewLimiter(float64(s.rate)))
    63  	}
    64  
    65  	if s.compression == "zstd" {
    66  		// We do not want to bog down the CPU with compression: we are fine with
    67  		// slowing down the snapshotting process, as it potentially may affect the
    68  		// WAL writes if they are located on the same disk.
    69  		var enc *zstd.Encoder
    70  		if enc, err = zstd.NewWriter(dst, zstd.WithEncoderConcurrency(1)); err != nil {
    71  			level.Error(s.logger).Log("msg", "failed to create zstd encoder", "err", err)
    72  			return err
    73  		}
    74  		defer func() {
    75  			if err = enc.Close(); err != nil {
    76  				level.Error(s.logger).Log("msg", "zstd compression failed", "err", err)
    77  			}
    78  		}()
    79  		// Wrap the writer with the encoder in the last turn:
    80  		// we want to apply the rate limiting to the sink, not
    81  		// the encoder.
    82  		dst = enc
    83  	}
    84  
    85  	level.Info(s.logger).Log("msg", "persisting snapshot", "compression", s.compression, "rate_limit_mb", s.rate)
    86  	if _, err = s.tx.WriteTo(dst); err != nil {
    87  		level.Error(s.logger).Log("msg", "failed to write snapshot", "err", err)
    88  		return err
    89  	}
    90  
    91  	s.metrics.boltDBPersistSnapshotSize.Set(float64(w.written))
    92  	return nil
    93  }
    94  
    95  type writeCounter struct {
    96  	io.Writer
    97  	written int64
    98  }
    99  
   100  func (w *writeCounter) Write(p []byte) (int, error) {
   101  	n, err := w.Writer.Write(p)
   102  	w.written += int64(n)
   103  	return n, err
   104  }
   105  
   106  func (s *snapshotWriter) Release() {
   107  	if s.tx != nil {
   108  		// This is an in-memory rollback, no error expected.
   109  		_ = s.tx.Rollback()
   110  	}
   111  }
   112  
   113  type snapshotReader struct {
   114  	io.ReadCloser
   115  }
   116  
   117  var zstdMagic = []byte{0x28, 0xB5, 0x2F, 0xFD}
   118  
   119  func newSnapshotReader(snapshot io.ReadCloser) (*snapshotReader, error) {
   120  	b := bufio.NewReader(snapshot)
   121  	magic, err := b.Peek(4)
   122  	if err != nil {
   123  		return nil, err
   124  	}
   125  
   126  	s := snapshotReader{ReadCloser: io.NopCloser(b)}
   127  	if bytes.Equal(magic, zstdMagic) {
   128  		var dec *zstd.Decoder
   129  		if dec, err = zstd.NewReader(b); err != nil {
   130  			return nil, err
   131  		}
   132  		s.ReadCloser = dec.IOReadCloser()
   133  	}
   134  
   135  	return &s, nil
   136  }