github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/util/log.go (about)

     1  package util
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	"github.com/rs/zerolog"
     9  )
    10  
    11  // LogProgressFunc is a function that can be called to add to the progress
    12  type LogProgressFunc func(addProgress int)
    13  
    14  type LogProgressConfig struct {
    15  	Message string
    16  	Total   int
    17  	Sampler zerolog.Sampler
    18  }
    19  
    20  func DefaultLogProgressConfig(
    21  	message string,
    22  	total int,
    23  ) LogProgressConfig {
    24  	nth := uint32(total / 10) // sample every 10% by default
    25  	if nth == 0 {
    26  		nth = 1
    27  	}
    28  
    29  	sampler := newProgressLogsSampler(nth, 60*time.Second)
    30  	return NewLogProgressConfig(
    31  		message,
    32  		total,
    33  		sampler,
    34  	)
    35  }
    36  
    37  func NewLogProgressConfig(
    38  	message string,
    39  	total int,
    40  	sampler zerolog.Sampler) LogProgressConfig {
    41  	return LogProgressConfig{
    42  		Message: message,
    43  		Total:   total,
    44  		Sampler: sampler,
    45  	}
    46  
    47  }
    48  
    49  type LogProgressOption func(config *LogProgressConfig)
    50  
    51  // LogProgress takes a total and return function such that when called adds the given
    52  // number to the progress and logs the progress every 10% or every 60 seconds whichever
    53  // comes first.
    54  // The returned function can be called concurrently.
    55  // An eta is also logged, but it assumes that the progress is linear.
    56  func LogProgress(
    57  	log zerolog.Logger,
    58  	config LogProgressConfig,
    59  ) LogProgressFunc {
    60  	sampler := log.Sample(config.Sampler)
    61  
    62  	start := time.Now()
    63  	currentIndex := uint64(0)
    64  	return func(add int) {
    65  		current := atomic.AddUint64(&currentIndex, uint64(add))
    66  
    67  		percentage := float64(100)
    68  		if config.Total > 0 {
    69  			percentage = (float64(current) / float64(config.Total)) * 100.
    70  		}
    71  		elapsed := time.Since(start)
    72  		elapsedString := elapsed.Round(1 * time.Second).String()
    73  
    74  		etaString := "unknown"
    75  		if percentage > 0 {
    76  			eta := time.Duration(float64(elapsed) / percentage * (100 - percentage))
    77  			if eta < 0 {
    78  				eta = 0
    79  			}
    80  			etaString = eta.Round(1 * time.Second).String()
    81  
    82  		}
    83  
    84  		if current != uint64(config.Total) {
    85  			sampler.Info().Msgf("%s progress %d/%d (%.1f%%) elapsed: %s, eta %s", config.Message, current, config.Total, percentage, elapsedString, etaString)
    86  		} else {
    87  			log.Info().Msgf("%s progress %d/%d (%.1f%%) total time %s", config.Message, current, config.Total, percentage, elapsedString)
    88  		}
    89  	}
    90  }
    91  
    92  type TimedSampler struct {
    93  	start    time.Time
    94  	Duration time.Duration
    95  	mu       sync.Mutex
    96  }
    97  
    98  var _ zerolog.Sampler = (*TimedSampler)(nil)
    99  
   100  func NewTimedSampler(duration time.Duration) *TimedSampler {
   101  	return &TimedSampler{
   102  		start:    time.Now(),
   103  		Duration: duration,
   104  		mu:       sync.Mutex{},
   105  	}
   106  }
   107  
   108  func (s *TimedSampler) Sample(_ zerolog.Level) bool {
   109  	s.mu.Lock()
   110  	defer s.mu.Unlock()
   111  
   112  	if time.Since(s.start) > s.Duration {
   113  		s.start = time.Now()
   114  		return true
   115  	}
   116  	return false
   117  }
   118  
   119  func (s *TimedSampler) Reset() {
   120  	s.mu.Lock()
   121  	defer s.mu.Unlock()
   122  
   123  	s.start = time.Now()
   124  }
   125  
   126  type progressLogsSampler struct {
   127  	basicSampler *zerolog.BasicSampler
   128  	timedSampler *TimedSampler
   129  }
   130  
   131  var _ zerolog.Sampler = (*progressLogsSampler)(nil)
   132  
   133  // newProgressLogsSampler returns a sampler that samples every nth log
   134  // and also samples a log if the last log was more than duration ago
   135  func newProgressLogsSampler(nth uint32, duration time.Duration) zerolog.Sampler {
   136  	return &progressLogsSampler{
   137  		basicSampler: &zerolog.BasicSampler{N: nth},
   138  		timedSampler: NewTimedSampler(duration),
   139  	}
   140  }
   141  
   142  func (s *progressLogsSampler) Sample(lvl zerolog.Level) bool {
   143  	sample := s.basicSampler.Sample(lvl)
   144  	if sample {
   145  		s.timedSampler.Reset()
   146  		return true
   147  	}
   148  	return s.timedSampler.Sample(lvl)
   149  }