Fix progress bar for "get" command

License: MIT
Signed-off-by: rht <[email protected]>

core/commands/get.go

@@ -71,7 +71,9 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 			// only use this when the flag is '-C' without '-a'
 			reader, err = getZip(req.Context().Context, node, p, cmplvl)
 		} else {
-			reader, err = get(req.Context().Context, node, p, cmplvl)
+			var length uint64
+			reader, length, err = get(req.Context().Context, node, p, cmplvl)
+			res.SetLength(length)
 		}
 		if err != nil {
 			res.SetError(err, cmds.ErrNormal)
@@ -114,6 +116,8 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 			}
 			defer file.Close()
 
+			// bar length is set to 0 because there is no deterministic way to
+			// pre-determine the size of the .tar.gz output
 			bar := pb.New(0).SetUnits(pb.U_BYTES)
 			bar.Output = os.Stderr
 			pbReader := bar.NewProxyReader(outReader)
@@ -131,8 +135,7 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 
 		fmt.Printf("Saving file(s) to %s\n", outPath)
 
-		// TODO: get total length of files
-		bar := pb.New(0).SetUnits(pb.U_BYTES)
+		bar := pb.New64(int64(res.Length())).SetUnits(pb.U_BYTES)
 		bar.Output = os.Stderr
 
 		// wrap the reader with the progress bar proxy reader
@@ -142,8 +145,7 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 		defer bar.Finish()
 
 		extractor := &tar.Extractor{outPath}
-		err = extractor.Extract(reader)
-		if err != nil {
+		if err = extractor.Extract(reader); err != nil {
 			res.SetError(err, cmds.ErrNormal)
 		}
 	},
@@ -163,13 +165,20 @@ func getCompressOptions(req cmds.Request) (int, error) {
 	return gzip.NoCompression, nil
 }
 
-func get(ctx context.Context, node *core.IpfsNode, p path.Path, compression int) (io.Reader, error) {
+func get(ctx context.Context, node *core.IpfsNode, p path.Path, compression int) (io.Reader, uint64, error) {
 	dagnode, err := core.Resolve(ctx, node, p)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
+	}
+
+	reader, err := utar.NewReader(p, node.DAG, dagnode, compression)
+
+	length, err := utar.GetTarSize(node.DAG, dagnode)
+	if err != nil {
+		return nil, 0, err
 	}
 
-	return utar.NewReader(p, node.DAG, dagnode, compression)
+	return reader, length, err
 }
 
 // getZip is equivalent to `ipfs getdag $hash | gzip`

test/sharness/t0090-get.sh

@@ -23,7 +23,7 @@ test_get_cmd() {
 	'
 
 	test_expect_success "ipfs get succeeds" '
-	  echo "Hello Worlds!" >data &&
+	  echo "Hello Worlds#" >data &&
 	  HASH=`ipfs add -q data` &&
 	  ipfs get "$HASH" >actual
 	'
@@ -80,7 +80,7 @@ test_get_cmd() {
 	  mkdir -p dir &&
 	  touch dir/a &&
 	  mkdir -p dir/b &&
-	  echo "Hello, Worlds!" >dir/b/c &&
+	  echo "Hello, Worlds#" >dir/b/c &&
 	  HASH2=`ipfs add -r -q dir | tail -n 1` &&
 	  ipfs get "$HASH2" >actual
 	'

unixfs/tar/reader.go

@@ -17,6 +17,8 @@ import (
 	upb "github.com/ipfs/go-ipfs/unixfs/pb"
 )
 
+const tarBlockSize = 512
+
 type Reader struct {
 	buf        bytes.Buffer
 	closed     bool
@@ -202,3 +204,43 @@ func (r *Reader) syncCopy(reader io.Reader) error {
 	}
 	return nil
 }
+
+func GetTarSize(dag mdag.DAGService, dagnode *mdag.Node) (uint64, error) {
+	return _GetTarSize(dag, dagnode, true)
+}
+
+func _GetTarSize(dag mdag.DAGService, dagnode *mdag.Node, isparent bool) (uint64, error) {
+	size := uint64(0)
+
+	pb := new(upb.Data)
+	if err := proto.Unmarshal(dagnode.Data, pb); err != nil {
+		return 0, err
+	}
+
+	if pb.GetType() == upb.Data_Directory {
+		//size += tarBlockSize
+
+		ctx, cancel := context.WithTimeout(context.TODO(), time.Second*60)
+		defer cancel()
+		for _, ng := range dag.GetDAG(ctx, dagnode) {
+			childNode, err := ng.Get(ctx)
+			if err != nil {
+				return 0, err
+			}
+			childSize, err := _GetTarSize(dag, childNode, false)
+			if err != nil {
+				return 0, err
+			}
+			size += childSize
+		}
+	} else {
+		unixSize := pb.GetFilesize()
+		// tar header + file size + round up to nearest 512 bytes
+		size = tarBlockSize + unixSize + (tarBlockSize - unixSize%tarBlockSize)
+	}
+
+	if isparent {
+		size += 2 * tarBlockSize // tar parent padding
+	}
+	return size, nil
+}