Make the command stop when finished

cmd/ru/main.go

@@ -44,26 +44,29 @@ func (u *UnitValue) Set(value string) error {
 func main() {
 	var path string
 	var depth int
-	var t tree.Node
+	var root tree.Node
 	var unit = &UnitValue{unit: "KB"}
 	flag.StringVar(&path, "path", ".", "Path from where to start the walk from")
 	flag.IntVar(&depth, "depth", 0, "Depth to display")
 	flag.Var(unit, "unit", "Unit in which to report size")
 	flag.Parse()
 
-	if depth == 0 {
+	if depth == 0 || !tree.AnyDirectoryDownThere(path) {
-		t = tree.NewSingle(path)
+		root = tree.NewSingle(path)
 	} else {
-		t = tree.NewTop(path)
+		root = tree.NewTop(path)
 	}
-	t.SetUnit(unit.String())
+	root.SetUnit(unit.String())
-	go t.Spawn(depth)
+	go root.Spawn(depth)
-	go t.Collect()
+
-	for !t.Complete() {
+	for {
 		select {
 		case <-time.After(500 * time.Millisecond):
 			fmt.Printf("\033[H\033[2J")
-			fmt.Println(t)
+			fmt.Println(root)
+			if root.Complete() {
+				return
+			}
 		}
 	}
 }

tree/dynamic_tree.go

@@ -14,17 +14,6 @@ const (
 	SingleType       = "single"
 )
 
-type ByteSize uint
-
-const (
-	B ByteSize = 1 << (10 * iota)
-	KB
-	MB
-	GB
-	TB
-	PB
-)
-
 var ErrUnknownType = errors.New("unknown type")
 var ErrShouldBeBottom = errors.New("should be of type Bottom")
 var ErrChannelError = errors.New("channel errored")
@@ -42,58 +31,10 @@ type Node interface {
 	Size() int64
 	Spawn(int) error
 	Collect() error
-	Collector() chan int64
+	AddCollector(chan int64)
 	Level() int
-}
+	Name() string
-
+	Complete() bool
-func isLastLevel(node Node, maxLevel int) bool {
-	return maxLevel-1 == node.Level()
-}
-
-func setUnit(unit string) ByteSize {
-	switch unit {
-	case "B":
-		return B
-	case "KB":
-		return KB
-	case "MB":
-		return MB
-	case "GB":
-		return GB
-	case "TB":
-		return TB
-	case "PB":
-		// I reckon this will never happen, but you never know
-		return PB
-	default:
-		panic(ErrUnknownUnit)
-	}
-}
-
-func fmtUnit(unit ByteSize) string {
-	switch unit {
-	case B:
-		return "B"
-	case KB:
-		return "KB"
-	case MB:
-		return "MB"
-	case GB:
-		return "GB"
-	case TB:
-		return "TB"
-	case PB:
-		// I reckon this will never happen, but you never know
-		return "PB"
-	default:
-		panic(ErrUnknownUnit)
-	}
-}
-
-func fmtSize(size int64, unit ByteSize) string {
-	var dimension float64
-	dimension = float64(size) / float64(unit)
-	return fmt.Sprintf("%.2f %s", dimension, fmtUnit(unit))
 }
 
 //////////////////////////////
@@ -101,17 +42,17 @@ func fmtSize(size int64, unit ByteSize) string {
 //////////////////////////////
 
 type Top struct {
-	path    string
+	path     string
-	size    int64
+	size     int64
-	unit    ByteSize
+	unit     ByteSize
-	tree    []Node
+	tree     []Node
-	collect chan int64
+	collect  []<-chan int64
-	level   int
+	level    int
+	complete bool
 }
 
 func NewTop(path string) *Top {
 	t := &Top{path: path, level: 0, unit: KB}
-	t.collect = make(chan int64)
 	return t
 }
 
@@ -140,32 +81,34 @@ func (t *Top) Spawn(maxLevel int) error {
 		switch mode := info.Mode(); {
 		case mode.IsDir():
 			var child Node
+			var collect chan int64
 			path := filepath.Join(t.path, info.Name())
 			if isLastLevel(t, maxLevel) || !AnyDirectoryDownThere(path) {
-				child = t.newBottom(path)
+				child, collect = t.newBottom(path)
 			} else {
-				child = t.newIntermediate(path)
+				child, collect = t.newIntermediate(path)
 			}
+			t.AddCollector(collect)
 			go child.Spawn(maxLevel)
-			defer func() {
-				go child.Collect()
-			}()
 		case !mode.IsDir() && mode.IsRegular():
-			t.collect <- info.Size()
+			t.size += info.Size()
 		}
 	}
+
+	go t.Collect()
 	return nil
 }
 
 func (t *Top) Collect() error {
-	for s := range t.collect {
+	for s := range merge(t.collect) {
 		t.size += s
 	}
+	t.complete = true
 	return nil
 }
 
-func (t *Top) Collector() chan int64 {
+func (t *Top) AddCollector(collect chan int64) {
-	return t.collect
+	t.collect = append(t.collect, collect)
 }
 
 func (t *Top) Level() int {
@@ -176,6 +119,15 @@ func (t *Top) Name() string {
 	return filepath.Base(t.path)
 }
 
+func (t *Top) Complete() bool {
+	for _, child := range t.tree {
+		if !child.Complete() {
+			return false
+		}
+	}
+	return true
+}
+
 func (t *Top) String() string {
 	var out string
 	var lines []string
@@ -203,27 +155,29 @@ func (t *Top) addToTree(child Node) {
 	t.tree = append(t.tree, child)
 }
 
-func (t *Top) newChild(path, kind string) Node {
+func (t *Top) newChild(path, kind string) (Node, chan int64) {
 	switch {
 	case kind == IntermediateType:
-		i := NewIntermediate(path, t)
+		i, collect := NewIntermediate(path, t)
 		t.addToTree(i)
-		return i
+		return i, collect
 	case kind == BottomType:
-		b := NewBottom(path, t)
+		b, collect := NewBottom(path, t)
 		t.addToTree(b)
-		return b
+		return b, collect
 	default:
 		panic(ErrUnknownType)
 	}
 }
 
-func (t *Top) newIntermediate(path string) *Intermediate {
+func (t *Top) newIntermediate(path string) (*Intermediate, chan int64) {
-	return t.newChild(path, IntermediateType).(*Intermediate)
+	child, collect := t.newChild(path, IntermediateType)
+	return child.(*Intermediate), collect
 }
 
-func (t *Top) newBottom(path string) *Bottom {
+func (t *Top) newBottom(path string) (*Bottom, chan int64) {
-	return t.newChild(path, BottomType).(*Bottom)
+	child, collect := t.newChild(path, BottomType)
+	return child.(*Bottom), collect
 }
 
 //////////////////////////////
@@ -231,22 +185,24 @@ func (t *Top) newBottom(path string) *Bottom {
 //////////////////////////////
 
 type Intermediate struct {
-	path    string
+	path     string
-	size    int64
+	size     int64
-	unit    ByteSize
+	unit     ByteSize
-	parent  Node
+	parent   Node
-	tree    []Node
+	tree     []Node
-	collect chan int64
+	collect  []<-chan int64
-	refer   chan int64
+	refer    chan int64
-	level   int
+	level    int
+	complete bool
 }
 
-func NewIntermediate(path string, parent Node) *Intermediate {
+func NewIntermediate(path string, parent Node) (*Intermediate, chan int64) {
-	i := &Intermediate{path: path, refer: parent.Collector(), parent: parent}
+	refer := make(chan int64)
-	i.collect = make(chan int64)
+	i := &Intermediate{path: path, refer: refer, parent: parent}
+	i.collect = make([]<-chan int64, 1)
 	i.level = parent.Level() + 1
 	i.unit = parent.GetUnit()
-	return i
+	return i, refer
 }
 
 func (i *Intermediate) SetUnit(unit string) {
@@ -275,35 +231,36 @@ func (i *Intermediate) Spawn(maxLevel int) error {
 		switch mode := info.Mode(); {
 		case mode.IsDir():
 			var child Node
+			var collect chan int64
 			path := filepath.Join(i.path, info.Name())
 			if isLastLevel(i, maxLevel) || !AnyDirectoryDownThere(path) {
-				child = i.newBottom(path)
+				child, collect = i.newBottom(path)
 			} else {
-				child = i.newIntermediate(path)
+				child, collect = i.newIntermediate(path)
 			}
+			i.AddCollector(collect)
 			go child.Spawn(maxLevel)
-			defer func() {
-				go child.Collect()
-			}()
 		case !mode.IsDir() && mode.IsRegular():
-			i.collect <- info.Size()
+			i.size += info.Size()
+			i.refer <- info.Size()
 		}
 	}
 
+	go i.Collect()
 	return nil
 }
 
 func (i *Intermediate) Collect() error {
-	for s := range i.collect {
+	for s := range merge(i.collect) {
 		i.size += s
 		i.refer <- s
 	}
-	close(i.refer)
+	i.complete = true
 	return nil
 }
 
-func (i *Intermediate) Collector() chan int64 {
+func (i *Intermediate) AddCollector(collect chan int64) {
-	return i.collect
+	i.collect = append(i.collect, collect)
 }
 
 func (i *Intermediate) Level() int {
@@ -314,6 +271,15 @@ func (i *Intermediate) Name() string {
 	return filepath.Base(i.path)
 }
 
+func (i *Intermediate) Complete() bool {
+	for _, child := range i.tree {
+		if !child.Complete() {
+			return false
+		}
+	}
+	return true
+}
+
 func (i *Intermediate) String() string {
 	var lines []string
 	out := fmt.Sprintf("(%s) %s\n", fmtSize(i.Size(), i.GetUnit()), i.Name())
@@ -348,27 +314,29 @@ func (i *Intermediate) addToTree(child Node) {
 	i.tree = append(i.tree, child)
 }
 
-func (i *Intermediate) newChild(path, kind string) Node {
+func (i *Intermediate) newChild(path, kind string) (Node, chan int64) {
 	switch {
 	case kind == IntermediateType:
-		c := NewIntermediate(path, i)
+		c, collect := NewIntermediate(path, i)
 		i.addToTree(c)
-		return c
+		return c, collect
 	case kind == BottomType:
-		b := NewBottom(path, i)
+		b, collect := NewBottom(path, i)
 		i.addToTree(b)
-		return b
+		return b, collect
 	default:
 		panic(ErrUnknownType)
 	}
 }
 
-func (i *Intermediate) newIntermediate(path string) *Intermediate {
+func (i *Intermediate) newIntermediate(path string) (*Intermediate, chan int64) {
-	return i.newChild(path, IntermediateType).(*Intermediate)
+	child, collect := i.newChild(path, IntermediateType)
+	return child.(*Intermediate), collect
 }
 
-func (i *Intermediate) newBottom(path string) *Bottom {
+func (i *Intermediate) newBottom(path string) (*Bottom, chan int64) {
-	return i.newChild(path, BottomType).(*Bottom)
+	child, collect := i.newChild(path, BottomType)
+	return child.(*Bottom), collect
 }
 
 //////////////////////////////
@@ -376,22 +344,23 @@ func (i *Intermediate) newBottom(path string) *Bottom {
 //////////////////////////////
 
 type Bottom struct {
-	path    string
+	path     string
-	size    int64
+	size     int64
-	unit    ByteSize
+	unit     ByteSize
-	parent  Node
+	parent   Node
-	walker  Walker
+	walker   Walker
-	collect chan int64
+	collect  chan int64
-	refer   chan int64
+	refer    chan int64
-	level   int
+	level    int
+	complete bool
 }
 
-func NewBottom(path string, parent Node) *Bottom {
+func NewBottom(path string, parent Node) (*Bottom, chan int64) {
-	b := &Bottom{path: path, refer: parent.Collector(), parent: parent, unit: parent.GetUnit()}
+	refer := make(chan int64)
-	b.collect = make(chan int64)
+	b := &Bottom{path: path, refer: refer, parent: parent, unit: parent.GetUnit()}
 	b.level = parent.Level() + 1
 	b.unit = parent.GetUnit()
-	return b
+	return b, refer
 }
 
 func (b *Bottom) SetUnit(unit string) {
@@ -407,8 +376,11 @@ func (b *Bottom) Size() int64 {
 }
 
 func (b *Bottom) Spawn(maxLevel int) error {
-	b.walker = NewRealWalker(b.path, b.collect)
+	collect := make(chan int64)
+	b.AddCollector(collect)
+	b.walker = NewRealWalker(b.path, collect)
 	go b.walker.Walk()
+	go b.Collect()
 	return nil
 }
 
@@ -417,12 +389,12 @@ func (b *Bottom) Collect() error {
 		b.size += s
 		b.refer <- s
 	}
-	close(b.refer)
+	b.complete = true
 	return nil
 }
 
-func (b *Bottom) Collector() chan int64 {
+func (b *Bottom) AddCollector(collect chan int64) {
-	return b.collect
+	b.collect = collect
 }
 
 func (b *Bottom) Level() int {
@@ -433,6 +405,10 @@ func (b *Bottom) Name() string {
 	return filepath.Base(b.path)
 }
 
+func (b *Bottom) Complete() bool {
+	return b.complete
+}
+
 func (b *Bottom) String() string {
 	return fmt.Sprintf("(%s) %s", fmtSize(b.Size(), b.GetUnit()), b.Name())
 }
@@ -442,10 +418,27 @@ func (b *Bottom) String() string {
 //////////////////////////////
 
 type Single struct {
-	path    string
+	path     string
-	size    int64
+	size     int64
-	walker  Walker
+	unit     ByteSize
-	collect chan int64
+	walker   Walker
+	collect  chan int64
+	complete bool
+}
+
+func NewSingle(path string) *Single {
+	s := &Single{path: path}
+	collect := make(chan int64)
+	s.AddCollector(collect)
+	return s
+}
+
+func (s *Single) SetUnit(unit string) {
+	s.unit = setUnit(unit)
+}
+
+func (s *Single) GetUnit() ByteSize {
+	return s.unit
 }
 
 func (s *Single) Size() int64 {
@@ -455,6 +448,7 @@ func (s *Single) Size() int64 {
 func (s *Single) Spawn(maxLevel int) error {
 	s.walker = NewRealWalker(s.path, s.collect)
 	go s.walker.Walk()
+	go s.Collect()
 	return nil
 }
 
@@ -462,21 +456,28 @@ func (s *Single) Collect() error {
 	for v := range s.collect {
 		s.size += v
 	}
+	s.complete = true
 	return nil
 }
 
-func (s *Single) Collector() chan int64 {
+func (s *Single) AddCollector(collect chan int64) {
-	return s.collect
+	s.collect = collect
 }
 
 func (s *Single) Level() int {
 	return 0
 }
 
-func NewSingle(path string) *Single {
+func (s *Single) Name() string {
-	s := &Single{path: path}
+	return filepath.Base(s.path)
-	s.collect = make(chan int64)
+}
-	return s
+
+func (s *Single) Complete() bool {
+	return s.complete
+}
+
+func (s *Single) String() string {
+	return fmt.Sprintf("(%s) %s\n", fmtSize(s.Size(), s.GetUnit()), s.Name())
 }
 
 //////////////////////////////

tree/real_walker.go

@@ -23,6 +23,7 @@ func (r *RealWalker) walkFunc(path string, info os.FileInfo, err error) error {
 
 func (r *RealWalker) Walk() {
 	filepath.Walk(r.path, r.walkFunc)
+	close(r.report)
 }
 
 func NewRealWalker(path string, report chan int64) *RealWalker {

tree/tree_utils.go

@@ -0,0 +1,93 @@
+package tree
+
+import (
+	"fmt"
+	"sync"
+)
+
+type ByteSize uint
+
+const (
+	B ByteSize = 1 << (10 * iota)
+	KB
+	MB
+	GB
+	TB
+	PB
+)
+
+func isLastLevel(node Node, maxLevel int) bool {
+	return maxLevel-1 == node.Level()
+}
+
+func setUnit(unit string) ByteSize {
+	switch unit {
+	case "B":
+		return B
+	case "KB":
+		return KB
+	case "MB":
+		return MB
+	case "GB":
+		return GB
+	case "TB":
+		return TB
+	case "PB":
+		// I reckon this will never happen, but you never know
+		return PB
+	default:
+		panic(ErrUnknownUnit)
+	}
+}
+
+func fmtUnit(unit ByteSize) string {
+	switch unit {
+	case B:
+		return "B"
+	case KB:
+		return "KB"
+	case MB:
+		return "MB"
+	case GB:
+		return "GB"
+	case TB:
+		return "TB"
+	case PB:
+		// I reckon this will never happen, but you never know
+		return "PB"
+	default:
+		panic(ErrUnknownUnit)
+	}
+}
+
+func fmtSize(size int64, unit ByteSize) string {
+	var dimension float64
+	dimension = float64(size) / float64(unit)
+	return fmt.Sprintf("%.2f %s", dimension, fmtUnit(unit))
+}
+
+func merge(cs []<-chan int64) <-chan int64 {
+	var wg sync.WaitGroup
+	out := make(chan int64)
+
+	// Start an output goroutine for each input channel in cs.  output
+	// copies values from c to out until c is closed, then calls wg.Done.
+	output := func(c <-chan int64) {
+		for n := range c {
+			out <- n
+		}
+		wg.Done()
+	}
+	wg.Add(len(cs))
+	for _, c := range cs {
+		go output(c)
+	}
+
+	// Start a goroutine to close out once all the output goroutines are
+	// done.  This must start after the wg.Add call.
+	go func() {
+		wg.Wait()
+		close(out)
+	}()
+	return out
+}