diff --git a/claimtrie/claimtrie.go b/claimtrie/claimtrie.go index 555fd2d005..f99a147ed9 100644 --- a/claimtrie/claimtrie.go +++ b/claimtrie/claimtrie.go @@ -140,11 +140,10 @@ func New(cfg config.Config) (*ClaimTrie, error) { } err = trie.SetRoot(hash) // keep this after IncrementHeightTo if err == merkletrie.ErrFullRebuildRequired { - // TODO: pass in the interrupt signal here: - ct.runFullTrieRebuild(nil) + ct.runFullTrieRebuild(nil, cfg.Interrupt) } - if !ct.MerkleHash().IsEqual(hash) { + if interruptRequested(cfg.Interrupt) || !ct.MerkleHash().IsEqual(hash) { ct.Close() return nil, errors.Errorf("unable to restore the claim hash to %s at height %d", hash.String(), previousHeight) } @@ -247,7 +246,7 @@ func (ct *ClaimTrie) AppendBlock() error { names = append(names, expirations...) names = removeDuplicates(names) - nhns := ct.makeNameHashNext(names, false) + nhns := ct.makeNameHashNext(names, false, nil) for nhn := range nhns { ct.merkleTrie.Update(nhn.Name, nhn.Hash, true) @@ -284,7 +283,7 @@ func (ct *ClaimTrie) updateTrieForHashForkIfNecessary() bool { } node.LogOnce(fmt.Sprintf("Rebuilding all trie nodes for the hash fork at %d...", ct.height)) - ct.runFullTrieRebuild(nil) + ct.runFullTrieRebuild(nil, nil) // I don't think it's safe to allow interrupt during fork return true } @@ -330,7 +329,7 @@ func (ct *ClaimTrie) ResetHeight(height int32) error { } err = ct.merkleTrie.SetRoot(hash) if err == merkletrie.ErrFullRebuildRequired { - ct.runFullTrieRebuild(names) + ct.runFullTrieRebuild(names, nil) } if !ct.MerkleHash().IsEqual(hash) { @@ -339,14 +338,14 @@ func (ct *ClaimTrie) ResetHeight(height int32) error { return nil } -func (ct *ClaimTrie) runFullTrieRebuild(names [][]byte) { +func (ct *ClaimTrie) runFullTrieRebuild(names [][]byte, interrupt <-chan struct{}) { var nhns chan NameHashNext if names == nil { node.LogOnce("Building the entire claim trie in RAM...") - nhns = ct.makeNameHashNext(nil, true) + nhns = ct.makeNameHashNext(nil, true, interrupt) } else { - nhns = ct.makeNameHashNext(names, false) + nhns = ct.makeNameHashNext(names, false, interrupt) } for nhn := range nhns { @@ -423,7 +422,17 @@ type NameHashNext struct { Next int32 } -func (ct *ClaimTrie) makeNameHashNext(names [][]byte, all bool) chan NameHashNext { +func interruptRequested(interrupted <-chan struct{}) bool { + select { + case <-interrupted: // should never block on nil + return true + default: + } + + return false +} + +func (ct *ClaimTrie) makeNameHashNext(names [][]byte, all bool, interrupt <-chan struct{}) chan NameHashNext { inputs := make(chan []byte, 512) outputs := make(chan NameHashNext, 512) @@ -448,6 +457,9 @@ func (ct *ClaimTrie) makeNameHashNext(names [][]byte, all bool) chan NameHashNex go func() { if all { ct.nodeManager.IterateNames(func(name []byte) bool { + if interruptRequested(interrupt) { + return false + } clone := make([]byte, len(name)) copy(clone, name) // iteration name buffer is reused on future loops inputs <- clone @@ -455,6 +467,9 @@ func (ct *ClaimTrie) makeNameHashNext(names [][]byte, all bool) chan NameHashNex }) } else { for _, name := range names { + if interruptRequested(interrupt) { + break + } inputs <- name } } diff --git a/claimtrie/claimtrie_test.go b/claimtrie/claimtrie_test.go index 10619a63d7..d93db22704 100644 --- a/claimtrie/claimtrie_test.go +++ b/claimtrie/claimtrie_test.go @@ -240,7 +240,7 @@ func TestRebuild(t *testing.T) { r.NotEqual(*merkletrie.EmptyTrieHash, *m) ct.merkleTrie = merkletrie.NewRamTrie() - ct.runFullTrieRebuild(nil) + ct.runFullTrieRebuild(nil, nil) m2 := ct.MerkleHash() r.NotNil(m2) @@ -432,7 +432,7 @@ func TestNormalizationRollbackFuzz(t *testing.T) { } } if j > 7 { - ct.runFullTrieRebuild(nil) + ct.runFullTrieRebuild(nil, nil) h := ct.MerkleHash() r.True(h.IsEqual(hashes[len(hashes)-1])) } diff --git a/claimtrie/config/config.go b/claimtrie/config/config.go index 73e2fb33e6..4920ca1733 100644 --- a/claimtrie/config/config.go +++ b/claimtrie/config/config.go @@ -40,6 +40,8 @@ type Config struct { NodeRepoPebble pebbleConfig TemporalRepoPebble pebbleConfig MerkleTrieRepoPebble pebbleConfig + + Interrupt <-chan struct{} } type pebbleConfig struct { diff --git a/server.go b/server.go index 6d1c20ec44..87200dd61f 100644 --- a/server.go +++ b/server.go @@ -2729,6 +2729,7 @@ func newServer(listenAddrs, agentBlacklist, agentWhitelist []string, claimTrieCfg := claimtrieconfig.DefaultConfig claimTrieCfg.DataDir = cfg.DataDir + claimTrieCfg.Interrupt = interrupt var ct *claimtrie.ClaimTrie