From 69da6a2a295a7e403c24d89d731fb259cd19d4ba Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Mon, 26 Apr 2021 03:50:29 -0700
Subject: [PATCH] feat: allow checkpointing to forks

Previously, `lotus sync checkpoint` would only checkpoint on the current
chain. Now, it can switch to a new fork.
---
 chain/checkpoint.go    | 27 ++++++++++++++++++++-------
 node/impl/full/sync.go |  2 +-
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/chain/checkpoint.go b/chain/checkpoint.go
index 8f99d73e4..4718b294d 100644
--- a/chain/checkpoint.go
+++ b/chain/checkpoint.go
@@ -1,6 +1,7 @@
 package chain
 
 import (
+	"context"
 	"encoding/json"
 
 	"github.com/filecoin-project/lotus/chain/types"
@@ -36,17 +37,20 @@ func loadCheckpoint(ds dtypes.MetadataDS) (types.TipSetKey, error) {
 	return tsk, err
 }
 
-func (syncer *Syncer) SetCheckpoint(tsk types.TipSetKey) error {
+func (syncer *Syncer) SyncCheckpoint(ctx context.Context, tsk types.TipSetKey) error {
 	if tsk == types.EmptyTSK {
 		return xerrors.Errorf("called with empty tsk")
 	}
 
-	syncer.checkptLk.Lock()
-	defer syncer.checkptLk.Unlock()
-
 	ts, err := syncer.ChainStore().LoadTipSet(tsk)
 	if err != nil {
-		return xerrors.Errorf("cannot find tipset: %w", err)
+		tss, err := syncer.Exchange.GetBlocks(ctx, tsk, 1)
+		if err != nil {
+			return xerrors.Errorf("failed to fetch tipset: %w", err)
+		} else if len(tss) != 1 {
+			return xerrors.Errorf("expected 1 tipset, got %d", len(tss))
+		}
+		ts = tss[0]
 	}
 
 	hts := syncer.ChainStore().GetHeaviestTipSet()
@@ -54,11 +58,18 @@ func (syncer *Syncer) SetCheckpoint(tsk types.TipSetKey) error {
 	if err != nil {
 		return xerrors.Errorf("cannot determine whether checkpoint tipset is in main-chain: %w", err)
 	}
-
 	if !hts.Equals(ts) && !anc {
-		return xerrors.Errorf("cannot mark tipset as checkpoint, since it isn't in the main-chain: %w", err)
+		if err := syncer.collectChain(ctx, ts, hts); err != nil {
+			return xerrors.Errorf("failed to collect chain for checkpoint: %w", err)
+		}
+		if err := syncer.ChainStore().SetHead(ts); err != nil {
+			return xerrors.Errorf("failed to set the chain head: %w", err)
+		}
 	}
 
+	syncer.checkptLk.Lock()
+	defer syncer.checkptLk.Unlock()
+
 	tskBytes, err := json.Marshal(tsk)
 	if err != nil {
 		return err
@@ -69,6 +80,8 @@ func (syncer *Syncer) SetCheckpoint(tsk types.TipSetKey) error {
 		return err
 	}
 
+	// TODO: This is racy. as there may be a concurrent sync in progress.
+	// The only real solution is to checkpoint inside the chainstore, not here.
 	syncer.checkpt = tsk
 
 	return nil
diff --git a/node/impl/full/sync.go b/node/impl/full/sync.go
index 1a088fb77..2c697483b 100644
--- a/node/impl/full/sync.go
+++ b/node/impl/full/sync.go
@@ -104,7 +104,7 @@ func (a *SyncAPI) SyncIncomingBlocks(ctx context.Context) (<-chan *types.BlockHe
 
 func (a *SyncAPI) SyncCheckpoint(ctx context.Context, tsk types.TipSetKey) error {
 	log.Warnf("Marking tipset %s as checkpoint", tsk)
-	return a.Syncer.SetCheckpoint(tsk)
+	return a.Syncer.SyncCheckpoint(ctx, tsk)
 }
 
 func (a *SyncAPI) SyncMarkBad(ctx context.Context, bcid cid.Cid) error {