Go言語でコンソールにプログレスバーを表示するライブラリ「pb」

cheggaaa/pb

package main

import (
	"github.com/cheggaaa/pb"
	"time"
)

func main() {
	count := 100000
	bar := pb.StartNew(count)
	for i := 0; i < count; i++ {
		bar.Increment()
		time.Sleep(time.Millisecond)
	}
	bar.FinishPrint("The End!")
}

exampleに現実的なサンプルとしてファイルコピーの進捗状況をプログレスバーで表示するスクリプトがある。

https://github.com/cheggaaa/pb/blob/master/example/copy/copy.go

読んで気がついたことをメモ。

ProgressBarはio.Writerとio.Readerのinterfaceを実装している。

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

AddメソッドはProgressbarのcurrentを増加させるのだけど、計算にatomicパッケージを使っている。

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(atomic.AddInt64(&pb.current, int64(add)))
}

io.MultiWriterに実際のファイルコピーのwriterとProgressBarのインスタンスを渡すことで、ファイルコピーの進捗状況をProgressBarで把握できるようにしている。

	// create multi writer
	writer := io.MultiWriter(dest, bar)

	// and copy
	io.Copy(writer, source)

プログレスバーの描画はgoroutine上で行われている。
currentの読み出しもatomicパッケージを使っている。

// Start print
func (pb *ProgressBar) Start() {
	pb.startTime = time.Now()
	if pb.Total == 0 {
		pb.ShowBar = false
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	go pb.writer()
}

func (pb *ProgressBar) writer() {
	var c, oc int64
	oc = -1
	for {
		if pb.isFinish {
			break
		}
		c = atomic.LoadInt64(&pb.current)
		if c != oc {
			pb.write(c)
			oc = c
		}
		time.Sleep(pb.RefreshRate)
	}
}

currentへの書き込みと読み出しが並行で処理されるのでatomicパッケージを使う必要があるのかな