Task: create disk calculation with go

to cancel long-wait goroutine, context package provide really nice abstraction.
Here is code using context to cancel loopDir, so in client side, we don’t need worry about goroutine leak.

var tokens = make(chan struct{}, 20)
var done = make(chan struct{})

const (
	interval = 500
	timeout  = 1300
)

func main() {
	flag.Parse()
	start := time.Now()
	foldername := flag.Args()
	// create ctx with timeout
	ctx, cancel := context.WithTimeout(context.Background(), timeout*time.Millisecond)
	defer cancel()
	out := make(chan int64)
	wg := &sync.WaitGroup{}
	tick := time.Tick(interval * time.Millisecond)
	wg.Add(1)
	go func() {
		for _, name := range foldername {
			go loopDir(ctx, name, out, wg)
		}
	}()
	go func() {
		wg.Wait()
		close(out)
	}()

	go func() {
		os.Stdin.Read(make([]byte, 1))
		close(done)
	}()

	var fileNumber, nByte int64
	for {
		select {
		case <-done:
			// when user hit key, exit the program
			for range out {
			}
			return
		case size, ok := <-out:
			if !ok {
				fmt.Printf("time.Since(start) = %+v\n", time.Since(start))
				panic("show me")
				print(fileNumber, nByte)
				return
			}
			fileNumber++
			nByte += size
		case <-tick:
			print(fileNumber, nByte)
		}
	}
}

func print(fileNumber, nByte int64) {
	fmt.Printf("count %d files, total size: %.1f KB\n", fileNumber, float64(nByte)/1e3)
}

func loopDir(ctx context.Context, foldername string, out chan<- int64, wg *sync.WaitGroup) {
	defer wg.Done()
	select {
	case <-ctx.Done():
	    //ctx will cancel after timeout
		return
	default:
		for _, v := range spanDir(foldername) {
			if v.IsDir() {
				subDir := filepath.Join(foldername, v.Name())
				tokens <- struct{}{}
				wg.Add(1)
				go loopDir(ctx, subDir, out, wg)
				<-tokens
			} else {
				out <- v.Size()
			}
		}
	}
}

func spanDir(foldername string) []os.FileInfo {
	subFolder, err := ioutil.ReadDir(foldername)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Reading folder with error: %v\n", err)
		return nil
	}
	return subFolder
}

var tokens = make(chan struct{}, 20)
var done = make(chan struct{})

const (
	interval = 500
)

func main() {
	flag.Parse()
	foldername := flag.Args()
	out := make(chan int64)
	wg := &sync.WaitGroup{}
	tick := time.Tick(interval * time.Millisecond)
	wg.Add(1)
	go func() {
		for _, name := range foldername {
			go loopDir(name, out, wg)
		}
	}()
	go func() {
		wg.Wait()
		close(out)
	}()

	go func() {
		os.Stdin.Read(make([]byte, 1))
		close(done)
	}()

	var fileNumber, nByte int64
	for {
		select {
		case <-done:
			for range out {
			}
			return
		case size, ok := <-out:
			if !ok {
				print(fileNumber, nByte)
				return
			}
			fileNumber++
			nByte += size
		case <-tick:
			print(fileNumber, nByte)
		}
	}
}

func print(fileNumber, nByte int64) {
	fmt.Printf("count %d files, total size: %.1f KB\n", fileNumber, float64(nByte)/1e3)

}

func cancel() bool {
	select {
	case <-done:
		return true
	default:
		return false
	}
}

func loopDir(foldername string, out chan<- int64, wg *sync.WaitGroup) {
	defer wg.Done()
	if cancel() {
		return
	}
	for _, v := range spanDir(foldername) {
		if v.IsDir() {
			subDir := filepath.Join(foldername, v.Name())
			tokens <- struct{}{}
			wg.Add(1)
			go loopDir(subDir, out, wg)
			<-tokens
		} else {
			out <- v.Size()
		}
	}
}

func spanDir(foldername string) []os.FileInfo {
	if cancel() {
		return nil
	}
	subFolder, err := ioutil.ReadDir(foldername)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Reading folder with error: %v\n", err)
		return nil
	}
	return subFolder
}

Back to posts