Refactor deal cleanup and cancellation and remove deal updates subscription raciness (#116)

* remove deal handler subscription race

* better deal cleanup and cancellation

* lint green

* changes as per review

* changes as per review

Author: aarshkshah1992

api/market.go

@@ -2,7 +2,6 @@ package api
 
 import (
 	"context"
-	"time"
 
 	"github.com/google/uuid"
 
@@ -29,6 +28,5 @@ type Market interface {
 
 // ProviderDealRejectionInfo is the information sent by the Storage Provider to the Client when it rejects a valid deal.
 type ProviderDealRejectionInfo struct {
-	Reason  string
-	Backoff time.Duration
+	Reason string
 }

db/db.go

@@ -3,11 +3,14 @@ package db
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"io/ioutil"
 	"path"
 	"runtime"
 )
 
+var ErrNotFound = errors.New("not found")
+
 type Scannable interface {
 	Scan(dest ...interface{}) error
 }

db/funds.go

@@ -45,7 +45,7 @@ func (f *FundsDB) Untag(ctx context.Context, dealUuid uuid.UUID) (abi.TokenAmoun
 	err := row.Scan(&collat.marshalled, &pubMsg.marshalled)
 	if err != nil {
 		if err == sql.ErrNoRows {
-			return abi.NewTokenAmount(0), nil
+			return abi.NewTokenAmount(0), ErrNotFound
 		}
 		return abi.NewTokenAmount(0), fmt.Errorf("getting untagged amount: %w", err)
 	}

db/funds_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"testing"
 
+	"golang.org/x/xerrors"
+
 	"github.com/filecoin-project/go-state-types/abi"
 	"github.com/google/uuid"
 
@@ -27,7 +29,7 @@ func TestFundsDB(t *testing.T) {
 
 	dealUUID := uuid.New()
 	amt, err := db.Untag(ctx, dealUUID)
-	req.NoError(err)
+	req.True(xerrors.Is(err, ErrNotFound))
 	req.Equal(int64(0), amt.Int64())
 
 	err = db.Tag(ctx, dealUUID, abi.NewTokenAmount(1111), abi.NewTokenAmount(2222))

db/storage.go

@@ -41,7 +41,7 @@ func (s *StorageDB) Untag(ctx context.Context, dealUuid uuid.UUID) (uint64, erro
 	err := row.Scan(&ps.marshalled)
 	if err != nil {
 		if err == sql.ErrNoRows {
-			return 0, nil
+			return 0, ErrNotFound
 		}
 		return 0, fmt.Errorf("getting untagged amount: %w", err)
 	}

db/storage_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"testing"
 
+	"golang.org/x/xerrors"
+
 	"github.com/google/uuid"
 
 	"github.com/stretchr/testify/require"
@@ -25,7 +27,7 @@ func TestStorageDB(t *testing.T) {
 
 	dealUUID := uuid.New()
 	amt, err := db.Untag(ctx, dealUUID)
-	req.NoError(err)
+	req.True(xerrors.Is(err, ErrNotFound))
 	req.Equal(uint64(0), amt)
 
 	err = db.Tag(ctx, dealUUID, 1111)

fundmanager/fundmanager.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"sync"
 
 	"github.com/filecoin-project/boost/db"
 	"github.com/filecoin-project/go-address"
@@ -45,8 +44,6 @@ type FundManager struct {
 	api fundManagerAPI
 	db  *db.FundsDB
 	cfg Config
-
-	lk sync.RWMutex
 }
 
 func New(cfg Config) func(api v1api.FullNode, sqldb *sql.DB) *FundManager {
@@ -74,9 +71,6 @@ func (m *FundManager) TagFunds(ctx context.Context, dealUuid uuid.UUID, proposal
 		return fmt.Errorf("getting publish deals message wallet balance: %w", err)
 	}
 
-	m.lk.Lock()
-	defer m.lk.Unlock()
-
 	// Check that the provider has enough funds in escrow to cover the
 	// collateral requirement for the deal
 	tagged, err := m.totalTagged(ctx)
@@ -112,9 +106,6 @@ func (m *FundManager) TagFunds(ctx context.Context, dealUuid uuid.UUID, proposal
 // TotalTagged returns the total funds tagged for specific deals for
 // collateral and publish storage deals message
 func (m *FundManager) TotalTagged(ctx context.Context) (*db.TotalTagged, error) {
-	m.lk.RLock()
-	defer m.lk.RUnlock()
-
 	return m.totalTagged(ctx)
 }
 
@@ -131,9 +122,6 @@ func (m *FundManager) totalTagged(ctx context.Context) (*db.TotalTagged, error)
 // It's called when it's no longer necessary to prevent the funds from being
 // used for a different deal (eg because the deal failed / was published)
 func (m *FundManager) UntagFunds(ctx context.Context, dealUuid uuid.UUID) error {
-	m.lk.Lock()
-	defer m.lk.Unlock()
-
 	untaggedAmt, err := m.db.Untag(ctx, dealUuid)
 	if err != nil {
 		return fmt.Errorf("persisting untag funds for deal to DB: %w", err)

gql/resolver.go

@@ -122,7 +122,7 @@ func (r *resolver) DealUpdate(ctx context.Context, args struct{ ID graphql.ID })
 
 	dealUpdatesSub, err := r.provider.SubscribeDealUpdates(dealUuid)
 	if err != nil {
-		if xerrors.Is(err, storagemarket.ErrDealHandlerFound) {
+		if xerrors.Is(err, storagemarket.ErrDealHandlerNotFound) {
 			close(net)
 			return net, nil
 		}
@@ -190,13 +190,13 @@ func (r *resolver) DealNew(ctx context.Context) (<-chan *dealResolver, error) {
 }
 
 // mutation: dealCancel(id): ID
-func (r *resolver) DealCancel(ctx context.Context, args struct{ ID graphql.ID }) (graphql.ID, error) {
+func (r *resolver) DealCancel(_ context.Context, args struct{ ID graphql.ID }) (graphql.ID, error) {
 	dealUuid, err := toUuid(args.ID)
 	if err != nil {
 		return args.ID, err
 	}
 
-	err = r.provider.CancelDeal(ctx, dealUuid)
+	err = r.provider.CancelDeal(dealUuid)
 	return args.ID, err
 }
 

storagemanager/storagemanager.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"sync"
 
 	"github.com/filecoin-project/boost/db"
 	"github.com/filecoin-project/go-state-types/abi"
@@ -21,8 +20,6 @@ type Config struct {
 type StorageManager struct {
 	db  *db.StorageDB
 	cfg Config
-
-	lk sync.RWMutex
 }
 
 func New(cfg Config) func(sqldb *sql.DB) *StorageManager {
@@ -36,9 +33,6 @@ func New(cfg Config) func(sqldb *sql.DB) *StorageManager {
 
 // Tag
 func (m *StorageManager) Tag(ctx context.Context, dealUuid uuid.UUID, pieceSize abi.PaddedPieceSize) error {
-	m.lk.Lock()
-	defer m.lk.Unlock()
-
 	// Get the total tagged storage, so that we know how much is available.
 	tagged, err := m.totalTagged(ctx)
 	if err != nil {
@@ -61,9 +55,6 @@ func (m *StorageManager) Tag(ctx context.Context, dealUuid uuid.UUID, pieceSize
 
 // Untag
 func (m *StorageManager) Untag(ctx context.Context, dealUuid uuid.UUID) error {
-	m.lk.Lock()
-	defer m.lk.Unlock()
-
 	pieceSize, err := m.db.Untag(ctx, dealUuid)
 	if err != nil {
 		return fmt.Errorf("persisting untag storage for deal to DB: %w", err)

storagemarket/deal_execution.go

@@ -2,6 +2,7 @@ package storagemarket
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -28,49 +29,44 @@ import (
 	"github.com/libp2p/go-libp2p-core/event"
 )
 
-func (p *Provider) doDeal(deal *types.ProviderDealState) {
+func (p *Provider) doDeal(deal *types.ProviderDealState, dh *dealHandler) {
 	// Set up pubsub for deal updates
-	bus := eventbus.NewBus()
-	pub, err := bus.Emitter(&types.ProviderDealState{}, eventbus.Stateful)
+	pub, err := dh.bus.Emitter(&types.ProviderDealState{}, eventbus.Stateful)
 	if err != nil {
-		err := fmt.Errorf("failed to create event emitter: %w", err)
+		err = fmt.Errorf("failed to create event emitter: %w", err)
+		p.cleanupDeal(deal)
 		p.failDeal(pub, deal, err)
 		return
 	}
 
-	// Create a context that can be cancelled for this deal if the user wants
-	// to cancel the deal early
-	ctx, stop := context.WithCancel(p.ctx)
-	defer stop()
-
-	stopped := make(chan struct{})
-	defer close(stopped)
-
-	// Keep track of the fields to subscribe to or cancel the deal
-	p.dealHandlers.track(&dealHandler{
-		dealUuid: deal.DealUuid,
-		ctx:      ctx,
-		stop:     stop,
-		stopped:  stopped,
-		bus:      bus,
-	})
-
 	// build in-memory state
 	fi, err := os.Stat(deal.InboundFilePath)
 	if err != nil {
 		err := fmt.Errorf("failed to stat output file: %w", err)
+		p.cleanupDeal(deal)
 		p.failDeal(pub, deal, err)
 		return
 	}
 	deal.NBytesReceived = fi.Size()
 
 	// Execute the deal synchronously
-	if err := p.execDeal(ctx, pub, deal); err != nil {
+	if err := p.execDealUptoAddPiece(dh.providerCtx, pub, deal, dh); err != nil {
+		p.cleanupDeal(deal)
 		p.failDeal(pub, deal, err)
+		return
 	}
+
+	// deal has been sent for sealing -> we can cleanup the deal state now and simply watch the deal on chain
+	// to wait for deal completion/slashing and update the state in DB accordingly.
+	p.cleanupDeal(deal)
+
+	// TODO
+	// Watch deal on chain and change state in DB and emit notifications.
+	// Given that cleanup deal above also gets rid of the deal handler, subscriptions to deal updates from here on
+	// will fail, we can look into it when we implement deal completion.
 }
 
-func (p *Provider) execDeal(ctx context.Context, pub event.Emitter, deal *types.ProviderDealState) error {
+func (p *Provider) execDealUptoAddPiece(ctx context.Context, pub event.Emitter, deal *types.ProviderDealState, dh *dealHandler) error {
 	// publish "new deal" event
 	p.fireEventDealNew(deal)
 
@@ -81,10 +77,13 @@ func (p *Provider) execDeal(ctx context.Context, pub event.Emitter, deal *types.
 
 	// Transfer Data
 	if deal.Checkpoint < dealcheckpoints.Transferred {
-		if err := p.transferAndVerify(ctx, pub, deal); err != nil {
+		if err := p.transferAndVerify(dh.transferCtx, pub, deal); err != nil {
+			dh.transferCancelled(nil)
 			return fmt.Errorf("failed data transfer: %w", err)
 		}
 	}
+	// transfer can no longer be cancelled
+	dh.transferCancelled(errors.New("transfer already complete"))
 
 	// Publish
 	if deal.Checkpoint <= dealcheckpoints.Published {
@@ -100,10 +99,6 @@ func (p *Provider) execDeal(ctx context.Context, pub event.Emitter, deal *types.
 		}
 	}
 
-	// Watch deal on chain and change state in DB and emit notifications.
-
-	// TODO: Clean up deal when it completes
-	//d.cleanupDeal()
 	return nil
 }
 
@@ -319,19 +314,11 @@ func (p *Provider) addPiece(ctx context.Context, pub event.Emitter, deal *types.
 
 	p.addDealLog(deal.DealUuid, "Deal handed off to sealer successfully")
 
-	// Now that the deal has been handed off to sealer, we should untag storage from the staging area.
-	err = p.storageManager.Untag(ctx, deal.DealUuid)
-	if err != nil {
-		log.Errorw("untagging storage", "uuid", deal.DealUuid, "err", err)
-		return err
-	}
-
 	return p.updateCheckpoint(ctx, pub, deal, dealcheckpoints.AddedPiece)
 }
 
 func (p *Provider) failDeal(pub event.Emitter, deal *types.ProviderDealState, err error) {
-	p.cleanupDeal(p.ctx, deal)
-
+	log.Errorw("deal failed", "id", deal.DealUuid, "err", err)
 	// Update state in DB with error
 	deal.Checkpoint = dealcheckpoints.Complete
 	if xerrors.Is(err, context.Canceled) {
@@ -350,29 +337,34 @@ func (p *Provider) failDeal(pub event.Emitter, deal *types.ProviderDealState, er
 	if pub != nil {
 		p.fireEventDealUpdate(pub, deal)
 	}
-
-	select {
-	case p.failedDealsChan <- failedDealReq{deal, err}:
-	case <-p.ctx.Done():
-	}
 }
 
-func (p *Provider) cleanupDeal(ctx context.Context, deal *types.ProviderDealState) {
+func (p *Provider) cleanupDeal(deal *types.ProviderDealState) {
+	// remove the temp file created for inbound deal data
 	_ = os.Remove(deal.InboundFilePath)
 
-	// clean up tagged funds
-	err := p.fundManager.UntagFunds(ctx, deal.DealUuid)
-	if err != nil {
-		log.Errorw("untagging funds", "id", deal.DealUuid, "err", err)
+	// close and clean up the deal handler
+	dh := p.getDealHandler(deal.DealUuid)
+	if dh != nil {
+		dh.transferCancelled(errors.New("deal cleaned up"))
+		dh.close()
+		p.delDealHandler(deal.DealUuid)
 	}
 
-	// clean up storage tag
-	err = p.storageManager.Untag(ctx, deal.DealUuid)
-	if err != nil {
-		log.Errorw("untagging storage", "id", deal.DealUuid, "err", err)
+	done := make(chan struct{}, 1)
+	// submit req to event loop to untag tagged funds and storage space
+	select {
+	case p.finishedDealsChan <- finishedDealsReq{deal: deal, done: done}:
+	case <-p.ctx.Done():
 	}
 
-	p.dealHandlers.del(deal.DealUuid)
+	// wait for event loop to finish cleanup and return before we return from here
+	// so caller is guaranteed that all resources associated with the deal have been cleanedup before
+	// taking further action.
+	select {
+	case <-done:
+	case <-p.ctx.Done():
+	}
 }
 
 func (p *Provider) fireEventDealNew(deal *types.ProviderDealState) {