more features, better locking

storage/sbmock/sbmock.go

@@ -2,10 +2,12 @@ package sbmock
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"math/rand"
 	"sync"
+	"time"
 
 	"github.com/filecoin-project/go-sectorbuilder"
 	"golang.org/x/xerrors"
@@ -23,6 +25,9 @@ type SBMock struct {
 	nextSectorID uint64
 	rateLimit    chan struct{}
 
+	preCommitDelay time.Duration
+	commitDelay    time.Duration
+
 	lk sync.Mutex
 }
 
@@ -35,8 +40,19 @@ func NewMockSectorBuilder(threads int, ssize uint64) *SBMock {
 	}
 }
 
+const (
+	statePacking = iota
+	statePreCommit
+	stateCommit
+)
+
 type sectorState struct {
 	pieces [][]byte
+	failed bool
+
+	state int
+
+	lk sync.Mutex
 }
 
 func (sb *SBMock) RateLimit() func() {
@@ -50,13 +66,16 @@ func (sb *SBMock) RateLimit() func() {
 
 func (sb *SBMock) AddPiece(size uint64, sectorId uint64, r io.Reader, existingPieces []uint64) (sectorbuilder.PublicPieceInfo, error) {
 	sb.lk.Lock()
-	defer sb.lk.Unlock()
-
 	ss, ok := sb.sectors[sectorId]
 	if !ok {
-		ss = &sectorState{}
+		ss = &sectorState{
+			state: statePacking,
+		}
 		sb.sectors[sectorId] = ss
 	}
+	sb.lk.Unlock()
+	ss.lk.Lock()
+	defer ss.lk.Unlock()
 
 	b, err := ioutil.ReadAll(r)
 	if err != nil {
@@ -76,7 +95,7 @@ func (sb *SBMock) SectorSize() uint64 {
 
 func (sb *SBMock) AcquireSectorId() (uint64, error) {
 	sb.lk.Lock()
-	sb.lk.Unlock()
+	defer sb.lk.Unlock()
 	id := sb.nextSectorID
 	sb.nextSectorID++
 	return id, nil
@@ -91,11 +110,16 @@ func (sb *SBMock) GenerateFallbackPoSt(sectorbuilder.SortedPublicSectorInfo, [se
 }
 
 func (sb *SBMock) SealPreCommit(sid uint64, ticket sectorbuilder.SealTicket, pieces []sectorbuilder.PublicPieceInfo) (sectorbuilder.RawSealPreCommitOutput, error) {
-	_, ok := sb.sectors[sid]
+	sb.lk.Lock()
+	ss, ok := sb.sectors[sid]
+	sb.lk.Unlock()
 	if !ok {
 		return sectorbuilder.RawSealPreCommitOutput{}, xerrors.Errorf("no sector with id %d in sectorbuilder", sid)
 	}
 
+	ss.lk.Lock()
+	defer ss.lk.Unlock()
+
 	ussize := sectorbuilder.UserBytesForSectorSize(sb.sectorSize)
 
 	// TODO: verify pieces in sinfo.pieces match passed in pieces
@@ -109,6 +133,14 @@ func (sb *SBMock) SealPreCommit(sid uint64, ticket sectorbuilder.SealTicket, pie
 		return sectorbuilder.RawSealPreCommitOutput{}, xerrors.Errorf("aggregated piece sizes don't match up: %d != %d", sum, ussize)
 	}
 
+	if ss.state != statePacking {
+		return sectorbuilder.RawSealPreCommitOutput{}, xerrors.Errorf("cannot call pre-seal on sector not in 'packing' state")
+	}
+
+	time.Sleep(sb.preCommitDelay)
+
+	ss.state = statePreCommit
+
 	return sectorbuilder.RawSealPreCommitOutput{
 		CommD: randComm(),
 		CommR: randComm(),
@@ -116,6 +148,25 @@ func (sb *SBMock) SealPreCommit(sid uint64, ticket sectorbuilder.SealTicket, pie
 }
 
 func (sb *SBMock) SealCommit(sid uint64, ticket sectorbuilder.SealTicket, seed sectorbuilder.SealSeed, pieces []sectorbuilder.PublicPieceInfo, precommit sectorbuilder.RawSealPreCommitOutput) ([]byte, error) {
+	sb.lk.Lock()
+	ss, ok := sb.sectors[sid]
+	sb.lk.Unlock()
+	if !ok {
+		return nil, xerrors.Errorf("no such sector %d", sid)
+	}
+	ss.lk.Lock()
+	defer ss.lk.Unlock()
+
+	if ss.failed {
+		return nil, xerrors.Errorf("[mock] cannot commit failed sector %d", sid)
+	}
+
+	if ss.state != statePreCommit {
+		return nil, xerrors.Errorf("cannot commit sector that has not been precommitted")
+	}
+
+	time.Sleep(sb.commitDelay)
+
 	buf := make([]byte, 32)
 	rand.Read(buf)
 	return buf, nil
@@ -136,3 +187,29 @@ func (sb *SBMock) AddWorker(context.Context, sectorbuilder.WorkerCfg) (<-chan se
 func (sb *SBMock) TaskDone(context.Context, uint64, sectorbuilder.SealRes) error {
 	panic("nyi")
 }
+
+// Test Instrumentation Methods
+
+func (sb *SBMock) FailSector(sid uint64) error {
+	sb.lk.Lock()
+	defer sb.lk.Unlock()
+	ss, ok := sb.sectors[sid]
+	if !ok {
+		return fmt.Errorf("no such sector in sectorbuilder")
+	}
+
+	ss.failed = true
+	return nil
+}
+
+func (sb *SBMock) SetPreCommitDelay(d time.Duration) {
+	sb.lk.Lock()
+	defer sb.lk.Unlock()
+	sb.preCommitDelay = d
+}
+
+func (sb *SBMock) SetCommitDelay(d time.Duration) {
+	sb.lk.Lock()
+	defer sb.lk.Unlock()
+	sb.commitDelay = d
+}