github.com/aclements/go-misc@v0.0.0-20240129233631-2f6ede80790c/benchmany/status.go (about)

     1  // Copyright 2015 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"fmt"
     9  	"math"
    10  	"os"
    11  	"time"
    12  
    13  	"github.com/aclements/go-moremath/fit"
    14  	"golang.org/x/crypto/ssh/terminal"
    15  )
    16  
    17  type StatusReporter struct {
    18  	update chan<- statusUpdate
    19  	done   chan bool
    20  }
    21  
    22  type statusUpdate struct {
    23  	progress float64
    24  	message  string
    25  }
    26  
    27  func NewStatusReporter() *StatusReporter {
    28  	if os.Getenv("TERM") == "dumb" || !terminal.IsTerminal(1) {
    29  		return &StatusReporter{}
    30  	}
    31  	update := make(chan statusUpdate)
    32  	sr := &StatusReporter{update: update}
    33  	go sr.loop(update)
    34  	return sr
    35  }
    36  
    37  func (sr *StatusReporter) Progress(msg string, frac float64) {
    38  	if sr.update != nil {
    39  		sr.update <- statusUpdate{message: msg, progress: frac}
    40  	}
    41  }
    42  
    43  func (sr *StatusReporter) Message(msg string) {
    44  	if sr.update == nil {
    45  		fmt.Println(msg)
    46  	} else {
    47  		sr.update <- statusUpdate{message: msg, progress: -1}
    48  	}
    49  }
    50  
    51  func (sr *StatusReporter) Stop() {
    52  	if sr.update != nil {
    53  		sr.done = make(chan bool)
    54  		close(sr.update)
    55  		<-sr.done
    56  		sr.update = nil
    57  	}
    58  }
    59  
    60  func (sr *StatusReporter) loop(updates <-chan statusUpdate) {
    61  	const resetLine = "\r\x1b[2K"
    62  	const wrapOff = "\x1b[?7l"
    63  	const wrapOn = "\x1b[?7h"
    64  
    65  	tick := time.NewTicker(time.Second / 4)
    66  	defer tick.Stop()
    67  
    68  	var end time.Time
    69  	t0 := time.Now()
    70  
    71  	var times, progress, weights []float64
    72  	var msg string
    73  	for {
    74  		select {
    75  		case update, ok := <-updates:
    76  			if !ok {
    77  				fmt.Print(resetLine)
    78  				close(sr.done)
    79  				return
    80  			}
    81  			if update.progress == -1 {
    82  				fmt.Print(resetLine)
    83  				fmt.Println(update.message)
    84  				break
    85  			}
    86  			now := float64(time.Now().Sub(t0))
    87  			times = append(times, float64(now))
    88  			progress = append(progress, update.progress)
    89  			weights = append(weights, 0)
    90  			msg = update.message
    91  
    92  			// Compute ETA using linear regression with
    93  			// exponentially decaying weights.
    94  			const halfLife = 150 * time.Second
    95  			for i, t := range times {
    96  				weights[i] = math.Exp(-1 / float64(halfLife) * (now - t))
    97  			}
    98  			reg := fit.PolynomialRegression(times, progress, weights, 1)
    99  			a, b := reg.Coefficients[0], reg.Coefficients[1]
   100  
   101  			// The intercept of a + b*x - 1 is the ending
   102  			// time.
   103  			if b == 0 {
   104  				end = time.Time{}
   105  			} else {
   106  				end = t0.Add(time.Duration((1 - a) / b))
   107  			}
   108  
   109  		case <-tick.C:
   110  		}
   111  
   112  		var eta string
   113  
   114  		if end.IsZero() {
   115  			eta = "unknown"
   116  		} else {
   117  			etaDur := end.Sub(time.Now())
   118  			// Trim off sub-second precision.
   119  			etaDur -= etaDur % time.Second
   120  			if etaDur <= 0 {
   121  				eta = "0s"
   122  			} else {
   123  				eta = etaDur.String()
   124  			}
   125  		}
   126  		if msg == "" {
   127  			eta = "ETA " + eta
   128  		} else {
   129  			eta = ", ETA " + eta
   130  		}
   131  		// TODO: This isn't quite right. If we hit the right
   132  		// edge of the terminal, it won't wrap, but the
   133  		// right-most character will be the *last* character
   134  		// in the string, since terminal keeps overwriting it.
   135  		fmt.Printf("%s%s%s%s%s", resetLine, wrapOff, msg, eta, wrapOn)
   136  	}
   137  }