From d00a7fa2c04ada007f74d828a14b122f8863975c Mon Sep 17 00:00:00 2001 From: Blallo Date: Sat, 7 Mar 2020 17:23:06 +0100 Subject: [PATCH] Make the command stop when finished --- cmd/ru/main.go | 21 ++-- tree/dynamic_tree.go | 289 ++++++++++++++++++++++--------------------- tree/real_walker.go | 1 + tree/tree_utils.go | 93 ++++++++++++++ 4 files changed, 251 insertions(+), 153 deletions(-) create mode 100644 tree/tree_utils.go diff --git a/cmd/ru/main.go b/cmd/ru/main.go index 33e07d5..30c36c3 100644 --- a/cmd/ru/main.go +++ b/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 { - t = tree.NewSingle(path) + if depth == 0 || !tree.AnyDirectoryDownThere(path) { + root = tree.NewSingle(path) } else { - t = tree.NewTop(path) + root = tree.NewTop(path) } - t.SetUnit(unit.String()) - go t.Spawn(depth) - go t.Collect() - for !t.Complete() { + root.SetUnit(unit.String()) + go root.Spawn(depth) + + for { select { case <-time.After(500 * time.Millisecond): fmt.Printf("\033[H\033[2J") - fmt.Println(t) + fmt.Println(root) + if root.Complete() { + return + } } } } diff --git a/tree/dynamic_tree.go b/tree/dynamic_tree.go index 5de1dee..352e8ea 100644 --- a/tree/dynamic_tree.go +++ b/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 -} - -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)) + Name() string + Complete() bool } ////////////////////////////// @@ -101,17 +42,17 @@ func fmtSize(size int64, unit ByteSize) string { ////////////////////////////// type Top struct { - path string - size int64 - unit ByteSize - tree []Node - collect chan int64 - level int + path string + size int64 + unit ByteSize + tree []Node + collect []<-chan int64 + 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 { - return t.collect +func (t *Top) AddCollector(collect chan int64) { + 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 { - return t.newChild(path, IntermediateType).(*Intermediate) +func (t *Top) newIntermediate(path string) (*Intermediate, chan int64) { + child, collect := t.newChild(path, IntermediateType) + return child.(*Intermediate), collect } -func (t *Top) newBottom(path string) *Bottom { - return t.newChild(path, BottomType).(*Bottom) +func (t *Top) newBottom(path string) (*Bottom, chan int64) { + 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 - size int64 - unit ByteSize - parent Node - tree []Node - collect chan int64 - refer chan int64 - level int + path string + size int64 + unit ByteSize + parent Node + tree []Node + collect []<-chan int64 + refer chan int64 + level int + complete bool } -func NewIntermediate(path string, parent Node) *Intermediate { - i := &Intermediate{path: path, refer: parent.Collector(), parent: parent} - i.collect = make(chan int64) +func NewIntermediate(path string, parent Node) (*Intermediate, chan int64) { + refer := 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 { - return i.collect +func (i *Intermediate) AddCollector(collect chan int64) { + 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 { - return i.newChild(path, IntermediateType).(*Intermediate) +func (i *Intermediate) newIntermediate(path string) (*Intermediate, chan int64) { + child, collect := i.newChild(path, IntermediateType) + return child.(*Intermediate), collect } -func (i *Intermediate) newBottom(path string) *Bottom { - return i.newChild(path, BottomType).(*Bottom) +func (i *Intermediate) newBottom(path string) (*Bottom, chan int64) { + 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 - size int64 - unit ByteSize - parent Node - walker Walker - collect chan int64 - refer chan int64 - level int + path string + size int64 + unit ByteSize + parent Node + walker Walker + collect chan int64 + refer chan int64 + level int + complete bool } -func NewBottom(path string, parent Node) *Bottom { - b := &Bottom{path: path, refer: parent.Collector(), parent: parent, unit: parent.GetUnit()} - b.collect = make(chan int64) +func NewBottom(path string, parent Node) (*Bottom, chan int64) { + refer := 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 { - return b.collect +func (b *Bottom) AddCollector(collect chan int64) { + 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 - size int64 - walker Walker - collect chan int64 + path string + size int64 + unit ByteSize + 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 { - return s.collect +func (s *Single) AddCollector(collect chan int64) { + s.collect = collect } func (s *Single) Level() int { return 0 } -func NewSingle(path string) *Single { - s := &Single{path: path} - s.collect = make(chan int64) - return s +func (s *Single) Name() string { + return filepath.Base(s.path) +} + +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()) } ////////////////////////////// diff --git a/tree/real_walker.go b/tree/real_walker.go index edcd8fc..41e313a 100644 --- a/tree/real_walker.go +++ b/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 { diff --git a/tree/tree_utils.go b/tree/tree_utils.go new file mode 100644 index 0000000..30e5d28 --- /dev/null +++ b/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 +}