sync: implement --order-by flag to order transfers - fixes #1205

fs/sync/pipe.go

@@ -1,12 +1,19 @@
 package sync
 
 import (
+	"container/heap"
 	"context"
+	"strings"
 	"sync"
 
+	"github.com/pkg/errors"
 	"github.com/rclone/rclone/fs"
+	"github.com/rclone/rclone/fs/fserrors"
 )
 
+// compare two items for order by
+type lessFn func(a, b fs.ObjectPair) bool
+
 // pipe provides an unbounded channel like experience
 //
 // Note unlike channels these aren't strictly ordered.
@@ -17,15 +24,58 @@ type pipe struct {
 	closed    bool
 	totalSize int64
 	stats     func(items int, totalSize int64)
+	less      lessFn
 }
 
-func newPipe(stats func(items int, totalSize int64), maxBacklog int) *pipe {
-	return &pipe{
+func newPipe(orderBy string, stats func(items int, totalSize int64), maxBacklog int) (*pipe, error) {
+	less, err := newLess(orderBy)
+	if err != nil {
+		return nil, fserrors.FatalError(err)
+	}
+	p := &pipe{
 		c:     make(chan struct{}, maxBacklog),
 		stats: stats,
+		less:  less,
 	}
+	if p.less != nil {
+		heap.Init(p)
+	}
+	return p, nil
 }
 
+// Len satisfy heap.Interface - must be called with lock held
+func (p *pipe) Len() int {
+	return len(p.queue)
+}
+
+// Len satisfy heap.Interface - must be called with lock held
+func (p *pipe) Less(i, j int) bool {
+	return p.less(p.queue[i], p.queue[j])
+}
+
+// Swap satisfy heap.Interface - must be called with lock held
+func (p *pipe) Swap(i, j int) {
+	p.queue[i], p.queue[j] = p.queue[j], p.queue[i]
+}
+
+// Push satisfy heap.Interface - must be called with lock held
+func (p *pipe) Push(item interface{}) {
+	p.queue = append(p.queue, item.(fs.ObjectPair))
+}
+
+// Pop satisfy heap.Interface - must be called with lock held
+func (p *pipe) Pop() interface{} {
+	old := p.queue
+	n := len(old)
+	item := old[n-1]
+	old[n-1] = fs.ObjectPair{} // avoid memory leak
+	p.queue = old[0 : n-1]
+	return item
+}
+
+// Check interface satisfied
+var _ heap.Interface = (*pipe)(nil)
+
 // Put an pair into the pipe
 //
 // It returns ok = false if the context was cancelled
@@ -36,7 +86,12 @@ func (p *pipe) Put(ctx context.Context, pair fs.ObjectPair) (ok bool) {
 		return false
 	}
 	p.mu.Lock()
-	p.queue = append(p.queue, pair)
+	if p.less == nil {
+		// no order-by
+		p.queue = append(p.queue, pair)
+	} else {
+		heap.Push(p, pair)
+	}
 	size := pair.Src.Size()
 	if size > 0 {
 		p.totalSize += size
@@ -68,10 +123,14 @@ func (p *pipe) Get(ctx context.Context) (pair fs.ObjectPair, ok bool) {
 		}
 	}
 	p.mu.Lock()
-	pair = p.queue[0]
-	p.queue[0].Src = nil
-	p.queue[0].Dst = nil
-	p.queue = p.queue[1:]
+	if p.less == nil {
+		// no order-by
+		pair = p.queue[0]
+		p.queue[0] = fs.ObjectPair{} // avoid memory leak
+		p.queue = p.queue[1:]
+	} else {
+		pair = heap.Pop(p).(fs.ObjectPair)
+	}
 	size := pair.Src.Size()
 	if size > 0 {
 		p.totalSize -= size
@@ -101,3 +160,49 @@ func (p *pipe) Close() {
 	p.closed = true
 	p.mu.Unlock()
 }
+
+// newLess returns a less function for the heap comparison or nil if
+// one is not required
+func newLess(orderBy string) (less lessFn, err error) {
+	if orderBy == "" {
+		return nil, nil
+	}
+	parts := strings.Split(strings.ToLower(orderBy), ",")
+	if len(parts) > 2 {
+		return nil, errors.Errorf("bad --order-by string %q", orderBy)
+	}
+	switch parts[0] {
+	case "name":
+		less = func(a, b fs.ObjectPair) bool {
+			return a.Src.Remote() < b.Src.Remote()
+		}
+	case "size":
+		less = func(a, b fs.ObjectPair) bool {
+			return a.Src.Size() < b.Src.Size()
+		}
+	case "modtime":
+		less = func(a, b fs.ObjectPair) bool {
+			ctx := context.Background()
+			return a.Src.ModTime(ctx).Before(b.Src.ModTime(ctx))
+		}
+	default:
+		return nil, errors.Errorf("unknown --order-by comparison %q", parts[0])
+	}
+	descending := false
+	if len(parts) > 1 {
+		switch parts[1] {
+		case "ascending", "asc":
+		case "descending", "desc":
+			descending = true
+		default:
+			return nil, errors.Errorf("unknown --order-by sort direction %q", parts[1])
+		}
+	}
+	if descending {
+		oldLess := less
+		less = func(a, b fs.ObjectPair) bool {
+			return !oldLess(a, b)
+		}
+	}
+	return less, nil
+}