github.com/aclements/go-misc@v0.0.0-20240129233631-2f6ede80790c/go-weave/amb/progress.go (about)

     1  // Copyright 2016 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 amb
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"log"
    11  	"os"
    12  	"sync"
    13  	"sync/atomic"
    14  	"time"
    15  )
    16  
    17  var count int64
    18  
    19  var progress struct {
    20  	printLock sync.Mutex
    21  	stop      chan struct{}
    22  	done      chan struct{}
    23  }
    24  
    25  const resetLine = "\r\x1b[2K"
    26  
    27  func startProgress() {
    28  	progress.stop = make(chan struct{})
    29  	progress.done = make(chan struct{})
    30  
    31  	go func() {
    32  		// Redirect process stdout and stderr.
    33  		//
    34  		// Alternatively, we could dup our pipes over stdout
    35  		// and stderr, but then we're in the way of any
    36  		// runtime debug output.
    37  		origStdout, origStderr := os.Stdout, os.Stderr
    38  		newStdoutR, newStdoutW, err := os.Pipe()
    39  		if err != nil {
    40  			log.Fatalf("failed to create stdout self-pipe: %v", err)
    41  		}
    42  		newStderrR, newStderrW, err := os.Pipe()
    43  		if err != nil {
    44  			log.Fatalf("failed to create stderr self-pipe: %v", err)
    45  		}
    46  
    47  		defer func() {
    48  			os.Stdout, os.Stderr = origStdout, origStderr
    49  			// Stop the feeder. It will close the write sides.
    50  			newStdoutR.Close()
    51  			newStderrR.Close()
    52  		}()
    53  		os.Stdout, os.Stderr = newStdoutW, newStderrW
    54  		go pipeFeeder(newStdoutR, origStdout, origStdout)
    55  		go pipeFeeder(newStderrR, origStderr, origStderr)
    56  
    57  		report := func(final bool) {
    58  			progress.printLock.Lock()
    59  			fmt.Fprintf(origStderr, "%s%d done", resetLine, atomic.LoadInt64(&count))
    60  			if final {
    61  				fmt.Fprintf(origStderr, "\n")
    62  			}
    63  			progress.printLock.Unlock()
    64  		}
    65  		ticker := time.NewTicker(100 * time.Millisecond)
    66  	loop:
    67  		for {
    68  			report(false)
    69  
    70  			select {
    71  			case <-ticker.C:
    72  			case <-progress.stop:
    73  				report(true)
    74  				break loop
    75  			}
    76  		}
    77  		ticker.Stop()
    78  		close(progress.done)
    79  	}()
    80  }
    81  
    82  func pipeFeeder(r, w, pstream *os.File) {
    83  	var buf [256]byte
    84  	bol := true
    85  	for {
    86  		n, err := r.Read(buf[:])
    87  		if n == 0 {
    88  			break
    89  		}
    90  		if bol {
    91  			bol = false
    92  			// Stop progress printing.
    93  			progress.printLock.Lock()
    94  			// Clear the progress line.
    95  			pstream.WriteString(resetLine)
    96  		}
    97  		// Print this message.
    98  		if n, err = w.Write(buf[:n]); err != nil {
    99  			panic(err)
   100  		}
   101  		if bytes.HasSuffix(buf[:n], []byte("\n")) {
   102  			// Resume progress printing.
   103  			progress.printLock.Unlock()
   104  			bol = true
   105  		}
   106  	}
   107  	w.Close()
   108  }
   109  
   110  func stopProgress() {
   111  	close(progress.stop)
   112  	<-progress.done
   113  }