core/commands: Added progress bars for 'add'
squash! core/commands: Added progress bars for 'add' Use vendored progress bar lib in 'add'
Matt Bell committed
Jan 21, 2015 at 18:38 UTC
c3ea164f64b24ee871e1c30f44938120afac2861
1 file changed
+143
-8
core/commands/add.go
+143
-8
@@ -5,7 +5,9 @@ import (
5
"errors"
6
"fmt"
7
"io"
8
+ "os"
9
"path"
10
+ "strings"
11
12
cmds "github.com/jbenet/go-ipfs/commands"
13
files "github.com/jbenet/go-ipfs/commands/files"
@@ -16,14 +18,22 @@ import (
18
pinning "github.com/jbenet/go-ipfs/pin"
19
ft "github.com/jbenet/go-ipfs/unixfs"
20
u "github.com/jbenet/go-ipfs/util"
21
+
22
+ "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/cheggaaa/pb"
23
)
24
25
// Error indicating the max depth has been exceded.
26
var ErrDepthLimitExceeded = fmt.Errorf("depth limit exceeded")
27
28
+// how many bytes of progress to wait before sending a progress update message
29
+const progressReaderIncrement = 1024 * 256
30
+
31
+const progressOptionName = "progress"
32
+
33
type AddedObject struct {
25
- Name string
26
- Hash string
34
+ Name string
35
+ Hash string `json:",omitempty"`
36
+ Bytes int64 `json:",omitempty"`
37
}
38
39
var AddCmd = &cmds.Command{
@@ -43,6 +53,26 @@ remains to be implemented.
53
Options: []cmds.Option{
54
cmds.OptionRecursivePath, // a builtin option that allows recursive paths (-r, --recursive)
55
cmds.BoolOption("quiet", "q", "Write minimal output"),
56
+ cmds.BoolOption(progressOptionName, "p", "Stream progress data"),
57
+ },
58
+ PreRun: func(req cmds.Request) error {
59
+ req.SetOption(progressOptionName, true)
60
+
61
+ sizeFile, ok := req.Files().(files.SizeFile)
62
+ if !ok {
63
+ // we don't need to error, the progress bar just won't know how big the files are
64
+ return nil
65
+ }
66
+
67
+ size, err := sizeFile.Size()
68
+ if err != nil {
69
+ // see comment above
70
+ return nil
71
+ }
72
+ log.Debugf("Total size of file being added: %v\n", size)
73
+ req.Values()["size"] = size
74
+
75
+ return nil
76
},
77
Run: func(req cmds.Request, res cmds.Response) {
78
n, err := req.Context().GetNode()
@@ -51,6 +81,8 @@ remains to be implemented.
81
return
82
}
83
84
+ progress, _, _ := req.Option(progressOptionName).Bool()
85
+
86
outChan := make(chan interface{})
87
res.SetOutput((<-chan interface{})(outChan))
88
@@ -63,13 +95,87 @@ remains to be implemented.
95
return
96
}
97
66
- _, err = addFile(n, file, outChan)
98
+ _, err = addFile(n, file, outChan, progress)
99
if err != nil {
100
return
101
}
102
}
103
}()
104
},
105
+ PostRun: func(res cmds.Response) {
106
+ outChan, ok := res.Output().(<-chan interface{})
107
+ if !ok {
108
+ res.SetError(u.ErrCast(), cmds.ErrNormal)
109
+ return
110
+ }
111
+
112
+ wrapperChan := make(chan interface{})
113
+ res.SetOutput((<-chan interface{})(wrapperChan))
114
+
115
+ size := int64(0)
116
+ s, found := res.Request().Values()["size"]
117
+ if found {
118
+ size = s.(int64)
119
+ }
120
+ showProgressBar := size >= progressBarMinSize
121
+
122
+ var bar *pb.ProgressBar
123
+ var terminalWidth int
124
+ if showProgressBar {
125
+ bar = pb.New64(size).SetUnits(pb.U_BYTES)
126
+ bar.ManualUpdate = true
127
+ bar.Start()
128
+
129
+ // the progress bar lib doesn't give us a way to get the width of the output,
130
+ // so as a hack we just use a callback to measure the output, then git rid of it
131
+ terminalWidth = 0
132
+ bar.Callback = func(line string) {
133
+ terminalWidth = len(line)
134
+ bar.Callback = nil
135
+ bar.Output = os.Stderr
136
+ log.Infof("terminal width: %v\n", terminalWidth)
137
+ }
138
+ bar.Update()
139
+ }
140
+
141
+ go func() {
142
+ lastFile := ""
143
+ var totalProgress, prevFiles, lastBytes int64
144
+
145
+ for out := range outChan {
146
+ output := out.(*AddedObject)
147
+ if len(output.Hash) > 0 {
148
+ if showProgressBar {
149
+ // clear progress bar line before we print "added x" output
150
+ fmt.Fprintf(os.Stderr, "\r%s\r", strings.Repeat(" ", terminalWidth))
151
+ }
152
+ wrapperChan <- output
153
+
154
+ } else {
155
+ log.Debugf("add progress: %v %v\n", output.Name, output.Bytes)
156
+
157
+ if !showProgressBar {
158
+ continue
159
+ }
160
+
161
+ if len(lastFile) == 0 {
162
+ lastFile = output.Name
163
+ }
164
+ if output.Name != lastFile || output.Bytes < lastBytes {
165
+ prevFiles += lastBytes
166
+ lastFile = output.Name
167
+ }
168
+ lastBytes = output.Bytes
169
+ delta := prevFiles + lastBytes - totalProgress
170
+ totalProgress = bar.Add64(delta)
171
+
172
+ bar.Update()
173
+ }
174
+ }
175
+
176
+ close(wrapperChan)
177
+ }()
178
+ },
179
Marshalers: cmds.MarshalerMap{
180
cmds.Text: func(res cmds.Response) (io.Reader, error) {
181
outChan, ok := res.Output().(<-chan interface{})
@@ -144,12 +250,19 @@ func addNode(n *core.IpfsNode, node *dag.Node) error {
250
return nil
251
}
252
147
-func addFile(n *core.IpfsNode, file files.File, out chan interface{}) (*dag.Node, error) {
253
+func addFile(n *core.IpfsNode, file files.File, out chan interface{}, progress bool) (*dag.Node, error) {
254
if file.IsDirectory() {
149
- return addDir(n, file, out)
255
+ return addDir(n, file, out, progress)
256
+ }
257
+
258
+ // if the progress flag was specified, wrap the file so that we can send
259
+ // progress updates to the client (over the output channel)
260
+ var reader io.Reader = file
261
+ if progress {
262
+ reader = &progressReader{file: file, out: out}
263
}
264
152
- dns, err := add(n, []io.Reader{file})
265
+ dns, err := add(n, []io.Reader{reader})
266
if err != nil {
267
return nil, err
268
}
@@ -161,7 +274,7 @@ func addFile(n *core.IpfsNode, file files.File, out chan interface{}) (*dag.Node
274
return dns[len(dns)-1], nil // last dag node is the file.
275
}
276
164
-func addDir(n *core.IpfsNode, dir files.File, out chan interface{}) (*dag.Node, error) {
277
+func addDir(n *core.IpfsNode, dir files.File, out chan interface{}, progress bool) (*dag.Node, error) {
278
log.Infof("adding directory: %s", dir.FileName())
279
280
tree := &dag.Node{Data: ft.FolderPBData()}
@@ -175,7 +288,7 @@ func addDir(n *core.IpfsNode, dir files.File, out chan interface{}) (*dag.Node,
288
break
289
}
290
178
- node, err := addFile(n, file, out)
291
+ node, err := addFile(n, file, out, progress)
292
if err != nil {
293
return nil, err
294
}
@@ -215,3 +328,25 @@ func outputDagnode(out chan interface{}, name string, dn *dag.Node) error {
328
329
return nil
330
}
331
+
332
+type progressReader struct {
333
+ file files.File
334
+ out chan interface{}
335
+ bytes int64
336
+ lastProgress int64
337
+}
338
+
339
+func (i *progressReader) Read(p []byte) (int, error) {
340
+ n, err := i.file.Read(p)
341
+
342
+ i.bytes += int64(n)
343
+ if i.bytes-i.lastProgress >= progressReaderIncrement || err == io.EOF {
344
+ i.lastProgress = i.bytes
345
+ i.out <- &AddedObject{
346
+ Name: i.file.FileName(),
347
+ Bytes: i.bytes,
348
+ }
349
+ }
350
+
351
+ return n, err
352
+}