Prioritize cancel requests

cmd/lhsm-plugin-az-core/remove.go

@@ -21,8 +21,8 @@ type RemoveOptions struct {
 	OpStartTime   time.Time
 }
 
-func Remove(o RemoveOptions) error {
-	ctx, _ := context.WithTimeout(context.Background(), time.Minute * 3)
+func Remove(ctx context.Context, o RemoveOptions) error {
+	ctx, _ = context.WithTimeout(ctx, time.Minute * 3)
 	p := azblob.NewPipeline(o.Credential, azblob.PipelineOptions{})
 	blobPath := path.Join(o.ExportPrefix, o.BlobName)
 	blobURL := azblob.NewContainerURL(*o.ContainerURL, p).NewBlockBlobURL(blobPath)

cmd/lhsm-plugin-az/mover.go

@@ -7,7 +7,6 @@ import (
 	"net/url"
 	"path"
 	"strings"
-	"sync"
 	"time"
 
 	core "github.com/wastore/lemur/cmd/lhsm-plugin-az-core"
@@ -25,38 +24,12 @@ import (
 	"github.com/wastore/lemur/dmplugin"
 )
 
-type cancelMap struct {
-	m map[string]context.CancelFunc
-	l sync.Mutex
-}
-
-func (c *cancelMap) add(key string, v context.CancelFunc) {
-	c.l.Lock()
-	defer c.l.Unlock()
-	c.m[key] = v
-}
-
-func (c *cancelMap) delete(key string) {
-	c.l.Lock()
-	defer c.l.Unlock()
-	delete(c.m, key)
-}
-
-func (c *cancelMap) cancelFunc(key string) (context.CancelFunc, bool) {
-	c.l.Lock()
-	defer c.l.Unlock()
-	f, ok := c.m[key]
-	return f, ok
-}
-
 // Mover supports archiving/restoring data to/from Azure Storage
 type Mover struct {
 	name                string
 	cred                azblob.Credential
 	httpClient          *http.Client
 	config              *archiveConfig
-	configLock	    sync.Mutex
-	actions             cancelMap //Actions in progress
 
 	//*Channels to interact wtih SAS Manager
 	getSAS              chan chan string
@@ -83,7 +56,6 @@ func AzMover(cfg *archiveConfig, creds azblob.Credential, archiveID uint32) *Mov
 		},
 		getSAS:             make(chan chan string),
 		forceSASRefresh:    make(chan time.Time),
-		actions:            cancelMap{m: make(map[string]context.CancelFunc)},
 	}
 }
 
@@ -216,7 +188,7 @@ func(m *Mover) refreshCredential(prevSASCtx time.Time) bool {
 	}	
 }
 
-func (m *Mover) Archive(action dmplugin.Action) error {
+func (m *Mover) Archive(ctx context.Context, action dmplugin.Action) error {
 	if util.ShouldLog(pipeline.LogDebug) {
 		util.Log(pipeline.LogDebug, fmt.Sprintf("%s id:%d archive %s %s", m.name, action.ID(), action.PrimaryPath(), action.UUID()))
 	} else {
@@ -263,14 +235,6 @@ func (m *Mover) Archive(action dmplugin.Action) error {
 		return err
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	m.configLock.Lock()
-	m.actions.add(action.PrimaryPath(),cancel)
-	m.configLock.Unlock()
-
-
 	total, err := core.Archive(ctx, core.ArchiveOptions{
 		ContainerURL:  m.config.ContainerURL(),
 		ResourceSAS:   sas,
@@ -295,9 +259,6 @@ func (m *Mover) Archive(action dmplugin.Action) error {
 		return err
 	}
 
-	m.configLock.Lock()
-	m.actions.delete(action.PrimaryPath())
-	m.configLock.Unlock()
 	if util.ShouldLog(pipeline.LogDebug) {
 		util.Log(pipeline.LogDebug, fmt.Sprintf("%s id:%d Archived %d bytes in %v from %s to %s/%s", m.name, action.ID(), total,
 		time.Since(start),
@@ -322,7 +283,7 @@ func (m *Mover) Archive(action dmplugin.Action) error {
 }
 
 // Restore fulfills an HSM Restore request
-func (m *Mover) Restore(action dmplugin.Action) error {
+func (m *Mover) Restore(ctx context.Context, action dmplugin.Action) error {
 	if util.ShouldLog(pipeline.LogDebug) {
 		util.Log(pipeline.LogDebug, fmt.Sprintf("%s id:%d restore %s %s", m.name, action.ID(), action.PrimaryPath(), action.UUID()))
 	} else {
@@ -351,12 +312,6 @@ func (m *Mover) Restore(action dmplugin.Action) error {
 	if err != nil {
 		return err
 	}
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	m.configLock.Lock()
-	m.actions.add(action.PrimaryPath(), cancel)
-	m.configLock.Unlock()
 
 	contentLen, err := core.Restore(ctx, core.RestoreOptions{
 		ContainerURL:    m.config.ContainerURL(),
@@ -380,9 +335,6 @@ func (m *Mover) Restore(action dmplugin.Action) error {
 		return err
 	}
 
-	m.configLock.Lock()
-	m.actions.delete(action.PrimaryPath())
-	m.configLock.Unlock()
 	if util.ShouldLog(pipeline.LogDebug) {
 		util.Log(pipeline.LogDebug, fmt.Sprintf("%s id:%d Restored %d bytes in %v from %s to %s", m.name, action.ID(), contentLen,
 		time.Since(start),
@@ -400,7 +352,7 @@ func (m *Mover) Restore(action dmplugin.Action) error {
 }
 
 // Remove fulfills an HSM Remove request
-func (m *Mover) Remove(action dmplugin.Action) error { 
+func (m *Mover) Remove(ctx context.Context, action dmplugin.Action) error { 
 	util.Log(pipeline.LogDebug, fmt.Sprintf("%s id:%d remove %s %s", m.name, action.ID(), action.PrimaryPath(), action.UUID()))
 	rate.Mark(1)
 	if action.UUID() == "" {
@@ -418,7 +370,7 @@ func (m *Mover) Remove(action dmplugin.Action) error {
 		return err
 	}
 
-	err = core.Remove(core.RemoveOptions{
+	err = core.Remove(ctx, core.RemoveOptions{
 		ContainerURL:   m.config.ContainerURL(),
 		ResourceSAS:   sas,
 		BlobName:      srcObj,
@@ -437,15 +389,3 @@ func (m *Mover) Remove(action dmplugin.Action) error {
 
 	return nil
 }
-
-
-func (m *Mover) Cancel(action dmplugin.Action) error {
-	util.Log(pipeline.LogDebug, fmt.Sprintf("%s id:%d Cancel %s %s", m.name, action.ID(), action.PrimaryPath(), action.UUID()))
-	m.configLock.Lock()
-	defer m.configLock.Unlock()
-	if cancel, ok := m.actions.cancelFunc(action.PrimaryPath()); ok {
-		cancel()
-		return nil
-	}
-	return errors.New("Could not find action with specified UUID")
-}

dmplugin/dmclient.go

@@ -24,7 +24,9 @@ import (
 
 type (
 	// ActionHandler is function that implements one of the commands
-	ActionHandler func(Action) error
+	ActionHandler func(context.Context, Action) error
+
+	actionFunc func()
 
 	// DataMoverClient is the data mover client to the HSM agent
 	DataMoverClient struct {
@@ -101,23 +103,19 @@ type (
 	// Archiver defines an interface for data movers capable of
 	// fulfilling Archive requests
 	Archiver interface {
-		Archive(Action) error
+		Archive(context.Context, Action) error
 	}
 
 	// Restorer defines an interface for data movers capable of
 	// fulfilling Restore requests
 	Restorer interface {
-		Restore(Action) error
+		Restore(context.Context, Action) error
 	}
 
 	// Remover defines an interface for data movers capable of
 	// fulfilling Remove requests
 	Remover interface {
-		Remove(Action) error
-	}
-
-	Canceler interface {
-		Cancel(Action) error
+		Remove(context.Context, Action) error
 	}
 )
 
@@ -277,9 +275,6 @@ func NewMover(plugin *Plugin, cli pb.DataMoverClient, config *Config) *DataMover
 	if remover, ok := config.Mover.(Remover); ok {
 		actions[pb.Command_REMOVE] = remover.Remove
 	}
-	if canceler, ok := config.Mover.(Canceler); ok {
-		actions[pb.Command_CANCEL] = canceler.Cancel
-	}
 
 	return &DataMoverClient{
 		plugin:    plugin,
@@ -337,23 +332,87 @@ func (dm *DataMoverClient) registerEndpoint(ctx context.Context) (*pb.Handle, er
 	return handle, nil
 }
 
-func (dm *DataMoverClient) processActions(ctx context.Context) chan *pb.ActionItem {
-	actions := make(chan *pb.ActionItem)
+func (dm *DataMoverClient) processActions(parentCtx context.Context) chan actionFunc {
+	actions := make(chan actionFunc)
+	ctx, cancel := context.WithCancel(parentCtx)
+
+	var processAction func(context.Context, *dmAction)
+	var queueAction func(*dmAction)
+	var cancelAction func(*dmAction)
+	var cancelMap sync.Map
+
+	maxTryCount := 3
+	if r := os.Getenv("COPYTOOL_RETRY_COUNT"); r != "" {
+		if v, err := strconv.Atoi(r); err == nil {
+			maxTryCount = v
+		}
+	}
+
+	queueAction = func(action *dmAction) {
+		childCtx, cancel := context.WithCancel(parentCtx)
+		cancelMap.Store(action.PrimaryPath(), cancel)
+		// Caller of queueAction is either processActions or handlers. Need to spawn
+		// a new goroutine below to not block the caller
+		go func() {
+			select {
+			case <-ctx.Done():
+				cancelMap.Delete(action.PrimaryPath())
+			case actions <- func() { processAction(childCtx, action) }:
+			}
+		}()
+	}
+
+	processAction = func(childCtx context.Context, action *dmAction) {
+
+		actionFn, err := dm.getActionHandler(action.item.Op)
+		if err != nil {
+			util.NewAzCopyLogSanitizer().SanitizeLogMessage(err.Error())
+			return
+		}
+
+		err = actionFn(childCtx, action)
+
+		if err != nil && util.ShouldRetry(err) && action.item.TryCount < int64(maxTryCount) {
+			action.item.TryCount += 1
+			queueAction(action) //Use parent context
+			util.NewAzCopyLogSanitizer().SanitizeLogMessage("Retrying: " + err.Error())
+			return
+		}
+
+		action.Finish(err)
+		cancelMap.Delete(action.PrimaryPath()) // Delete from map
+	}
+
+	cancelAction = func(action *dmAction) {
+		// lookup in the map and cancel the context
+		cancel, ok := cancelMap.Load(action.item.PrimaryPath)
+		if !ok {
+			msg := fmt.Sprintf("Received cancel for a non-existent action: %s", action.item.PrimaryPath)
+			alert.Warnf(msg)
+			action.Finish(errors.New(msg))
+			return
+		}
+
+		alert.Writer().Log(fmt.Sprintf("id:%d Cancel %s", action.item.Id, action.item.PrimaryPath))
+		cancel.(context.CancelFunc)()
+		action.Finish(nil)
+	}
 
 	go func() {
+		defer cancel()
 		defer close(actions)
-		handle, ok := getHandle(ctx)
+		handle, ok := getHandle(parentCtx)
 		if !ok {
 			alert.Warn(errors.New("No context"))
 			return
 		}
-		stream, err := dm.rpcClient.GetActions(ctx, handle)
+		stream, err := dm.rpcClient.GetActions(parentCtx, handle)
 		if err != nil {
 			alert.Warn(errors.Wrap(err, "GetActions() failed"))
 			return
 		}
 		for {
-			action, err := stream.Recv()
+			item, err := stream.Recv()
 			if err != nil {
 				if err == io.EOF {
 					debug.Print("Shutting down dmclient action stream")
@@ -364,8 +423,17 @@ func (dm *DataMoverClient) processActions(ctx context.Context) chan *pb.ActionIt
 			}
 			// debug.Printf("Got message id:%d op: %v %v", action.Id, action.Op, action.PrimaryPath)
 
-			action.TryCount = 0 //first try
-			actions <- action
+			item.TryCount = 0 //first try
+			action := &dmAction{
+				status: dm.status,
+				item: item,
+			}
+
+			if (item.Op == pb.Command_CANCEL) {
+				cancelAction(action)
+			} else {
+				queueAction(action)
+			}
 		}
 
 	}()
@@ -408,40 +476,9 @@ func (dm *DataMoverClient) getActionHandler(op pb.Command) (ActionHandler, error
 	return fn, nil
 }
 
-func (dm *DataMoverClient) requeueItem(item *pb.ActionItem, actions chan *pb.ActionItem) {
-	item.TryCount += 1
-	debug.Printf("Retrying action %d.Trycount: %d", item.Id, item.TryCount)
-	actions <- item
-}
-
-func (dm *DataMoverClient) handler(name string, actions chan *pb.ActionItem) {
-	maxTryCount := 3
-	if r := os.Getenv("COPYTOOL_RETRY_COUNT"); r != "" {
-		if v, err := strconv.Atoi(r); err == nil {
-			maxTryCount = v
-		}
-	}
-
-	for item := range actions {
-		action := &dmAction{
-			status: dm.status,
-			item:   item,
-		}
-
-		actionFn, err := dm.getActionHandler(item.Op)
-		if err == nil {
-			err = actionFn(action)
-		}
-		// debug.Printf("completed (action: %v) %v ", action, ret)
-		if util.ShouldRetry(err) && item.TryCount < int64(maxTryCount) {
-			go dm.requeueItem(item, actions)
-		} else {
-			action.Finish(err)
-		}
-
-		if err != nil {
-			err = errors.New(util.NewAzCopyLogSanitizer().SanitizeLogMessage(err.Error()))
-		}
+func (dm *DataMoverClient) handler(name string, actions chan actionFunc) {
+	for action := range actions {
+		action()
 	}
 	debug.Printf("%s: stopping", name)
 }