Migrate renterd subscriber to coreutils types

chain/subscriber.go

@@ -30,6 +30,7 @@ type (
 	ChainStore interface {
 		ApplyChainUpdate(ctx context.Context, cu *Update) error
 		ChainIndex() (types.ChainIndex, error)
+		WalletStateElements(ctx context.Context) ([]types.StateElement, error)
 	}
 
 	ContractStore interface {
@@ -51,9 +52,10 @@ type (
 		syncSig         chan struct{}
 		csUnsubscribeFn func()
 
-		mu             sync.Mutex
-		closedChan     chan struct{}
-		knownContracts map[types.FileContractID]api.ContractState
+		mu                 sync.Mutex
+		closedChan         chan struct{}
+		knownContracts     map[types.FileContractID]api.ContractState
+		knownStateElements map[types.Hash256]types.StateElement
 	}
 )
 
@@ -73,14 +75,28 @@ func NewSubscriber(cm ChainManager, cs ChainStore, contracts ContractStore, wall
 
 		syncSig: make(chan struct{}, 1),
 
-		closedChan: make(chan struct{}),
+		closedChan:         make(chan struct{}),
+		knownStateElements: make(map[types.Hash256]types.StateElement),
 	}
 
+	// make sure we don't hang
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	defer cancel()
+
 	// subscribe ourselves to receive new contract ids
-	subscriber.knownContracts, subscriber.csUnsubscribeFn, err = contracts.AddContractStoreSubscriber(context.Background(), subscriber)
+	subscriber.knownContracts, subscriber.csUnsubscribeFn, err = contracts.AddContractStoreSubscriber(ctx, subscriber)
+	if err != nil {
+		return nil, err
+	}
+
+	// fetch all state elements from the database
+	elements, err := cs.WalletStateElements(ctx)
 	if err != nil {
 		return nil, err
 	}
+	for _, el := range elements {
+		subscriber.knownStateElements[types.Hash256(el.ID)] = el
+	}
 
 	return subscriber, nil
 }
@@ -377,12 +393,12 @@ func (cs *Subscriber) sync(index types.ChainIndex) error {
 			return fmt.Errorf("failed to fetch updates: %w", err)
 		}
 
-		// aggregate updates into one chain update
-		cu := NewChainUpdate()
-
 		// lock the subscriber while we apply the updates
 		cs.mu.Lock()
 
+		// aggregate updates into one chain update
+		cu := NewChainUpdate(cs.knownStateElements)
+
 		// revert chain updates
 		for _, cru := range crus {
 			cs.revertContractUpdate(cu, cru)

chain/update.go

@@ -11,12 +11,13 @@ import (
 
 type (
 	Update struct {
-		Index types.ChainIndex
-
+		Index               types.ChainIndex
+		StateElements       map[types.Hash256]types.StateElement
 		ContractUpdates     map[types.FileContractID]*ContractUpdate
 		HostUpdates         map[types.PublicKey]HostUpdate
 		WalletOutputUpdates map[types.Hash256]WalletOutputUpdate
 		WalletEventUpdates  []WalletEventUpdate
+		RevertIndices       []types.ChainIndex
 	}
 
 	ContractUpdate struct {
@@ -48,8 +49,9 @@ type (
 )
 
 // NewChainUpdate returns a new ChainUpdate.
-func NewChainUpdate() *Update {
+func NewChainUpdate(elements map[types.Hash256]types.StateElement) *Update {
 	return &Update{
+		StateElements:       elements,
 		ContractUpdates:     make(map[types.FileContractID]*ContractUpdate),
 		HostUpdates:         make(map[types.PublicKey]HostUpdate),
 		WalletOutputUpdates: make(map[types.Hash256]WalletOutputUpdate),
@@ -105,12 +107,17 @@ func (cu *Update) RemoveSiacoinElements(ids []types.SiacoinOutputID) error {
 // WalletStateElements returns all state elements in the database. It is used
 // to update the proofs of all state elements affected by the update.
 func (cu *Update) WalletStateElements() (elements []types.StateElement, _ error) {
+	// copy the output updates to the state
 	for id, el := range cu.WalletOutputUpdates {
-		elements = append(elements, types.StateElement{
+		cu.StateElements[id] = types.StateElement{
 			ID:          id,
 			LeafIndex:   el.Element.LeafIndex,
 			MerkleProof: el.Element.MerkleProof,
-		})
+		}
+	}
+	// loop over the state
+	for _, se := range cu.StateElements {
+		elements = append(elements, se)
 	}
 	return
 }
@@ -146,5 +153,7 @@ func (cu *Update) RevertIndex(index types.ChainIndex) error {
 		}
 	}
 
+	// add the index to the list of reverted indices
+	cu.RevertIndices = append(cu.RevertIndices, index)
 	return nil
 }

stores/chain.go

@@ -48,6 +48,13 @@ func (s *SQLStore) ApplyChainUpdate(ctx context.Context, cu *chain.Update) error
 			}
 		}
 
+		// process reverts
+		for _, reverted := range cu.RevertIndices {
+			if err := revertIndex(tx, reverted); err != nil {
+				return fmt.Errorf("%w; failed to revert index %d", err, reverted.Height)
+			}
+		}
+
 		// update chain index
 		if err := updateChainIndex(tx, cu.Index); err != nil {
 			return fmt.Errorf("%w; failed to update chain index", err)
@@ -57,6 +64,33 @@ func (s *SQLStore) ApplyChainUpdate(ctx context.Context, cu *chain.Update) error
 	})
 }
 
+// WalletStateElements implements the ChainStore interface and returns all state
+// elements in the database.
+func (s *SQLStore) WalletStateElements(ctx context.Context) ([]types.StateElement, error) {
+	type row struct {
+		ID          hash256
+		LeafIndex   uint64
+		MerkleProof merkleProof
+	}
+	var rows []row
+	if err := s.db.WithContext(ctx).
+		Model(&dbWalletOutput{}).
+		Select("output_id AS id", "leaf_index", "merkle_proof").
+		Find(&rows).
+		Error; err != nil {
+		return nil, err
+	}
+	elements := make([]types.StateElement, 0, len(rows))
+	for _, r := range rows {
+		elements = append(elements, types.StateElement{
+			ID:          types.Hash256(r.ID),
+			LeafIndex:   r.LeafIndex,
+			MerkleProof: r.MerkleProof.proof,
+		})
+	}
+	return elements, nil
+}
+
 func updateHosts(tx *gorm.DB, ann map[types.PublicKey]chain.HostUpdate) error {
 	if len(ann) == 0 {
 		return nil
@@ -207,3 +241,10 @@ func updateWalletEvent(tx *gorm.DB, weu chain.WalletEventUpdate) error {
 		Delete(&dbWalletEvent{}).
 		Error
 }
+
+func revertIndex(tx *gorm.DB, index types.ChainIndex) error {
+	if err := tx.Model(&dbWalletEvent{}).Where("height = ? AND block_id = ?", index.Height, hash256(index.ID)).Delete(&dbWalletEvent{}).Error; err != nil {
+		return err
+	}
+	return tx.Model(&dbWalletOutput{}).Where("height = ? AND block_id = ?", index.Height, hash256(index.ID)).Delete(&dbWalletOutput{}).Error
+}

stores/chain_test.go

@@ -20,7 +20,7 @@ func TestApplyChainUpdate(t *testing.T) {
 	// prepare chain update
 	hk := types.PublicKey{1}
 	want := types.ChainIndex{Height: 1, ID: types.BlockID{1}}
-	cu := chain.NewChainUpdate()
+	cu := chain.NewChainUpdate(nil)
 	cu.HostUpdates[hk] = chain.HostUpdate{
 		Announcement: chain.HostAnnouncement{NetAddress: "foo.com:1000"},
 		BlockHeight:  want.Height,
@@ -55,7 +55,7 @@ func TestApplyChainUpdate(t *testing.T) {
 	}
 
 	// update the contract
-	cu = chain.NewChainUpdate()
+	cu = chain.NewChainUpdate(nil)
 	cu.Index = types.ChainIndex{Height: 2, ID: types.BlockID{2}}
 	cu.ContractUpdates[types.FileContractID{1}] = &chain.ContractUpdate{State: api.ContractStateActive}
 	err = db.ApplyChainUpdate(context.Background(), cu)
@@ -72,7 +72,7 @@ func TestApplyChainUpdate(t *testing.T) {
 	}
 
 	// add wallet output & event
-	cu = chain.NewChainUpdate()
+	cu = chain.NewChainUpdate(nil)
 	cu.Index = types.ChainIndex{Height: 3, ID: types.BlockID{3}}
 	cu.WalletOutputUpdates[types.Hash256{1}] = chain.WalletOutputUpdate{Addition: true, Element: wallet.SiacoinElement{
 		SiacoinElement: types.SiacoinElement{