diff --git a/backup/run.go b/backup/run.go index f36f951685a..00411c7f56f 100644 --- a/backup/run.go +++ b/backup/run.go @@ -240,7 +240,9 @@ func runRestoreCmd() error { } ctx, cancel := context.WithTimeout(context.Background(), time.Minute) defer cancel() - if _, err = zc.AssignIds(ctx, &pb.Num{Val: val, Type: typ}); err != nil { + // Bump rather than lease val more ids, as the online restore does: the lease must reach + // the backup's max, and a repeated restore must not move it again. + if err := worker.BumpLease(ctx, zc, val, typ); err != nil { fmt.Printf("Failed to assign %s %d in Zero: %v\n", pb.NumLeaseType_name[int32(typ)], val, err) return err diff --git a/worker/lease_test.go b/worker/lease_test.go new file mode 100644 index 00000000000..74be13a834d --- /dev/null +++ b/worker/lease_test.go @@ -0,0 +1,59 @@ +//go:build integration + +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package worker + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + + "github.com/dgraph-io/dgraph/v25/protos/pb" + "github.com/dgraph-io/dgraph/v25/testutil" +) + +// TestBumpLeaseIsIdempotent checks that bumping a lease raises it to the asked value once, and that +// bumping again to the same or a lower value leaves it where it is. A restore calls BumpLease with +// the backup's max uid, so restoring the same backup twice must not consume more ids. +func TestBumpLeaseIsIdempotent(t *testing.T) { + conn, err := grpc.NewClient(testutil.GetSockAddrZero(), + grpc.WithTransportCredentials(insecure.NewCredentials())) + require.NoError(t, err) + defer func() { require.NoError(t, conn.Close()) }() + zc := pb.NewZeroClient(conn) + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + // nextUid leases a single uid, which tells us where the lease currently stands. + nextUid := func() uint64 { + res, err := zc.AssignIds(ctx, &pb.Num{Val: 1, Type: pb.Num_UID}) + require.NoError(t, err) + return res.GetStartId() + } + + // Anything else in the cluster may lease uids while this runs, so the checks allow some slack + // rather than expecting exact values. Leasing val more ids instead of bumping would move the + // lease by at least target on every call, far past that slack. + const slack = 100 + target := nextUid() + 1000 + require.NoError(t, BumpLease(ctx, zc, target, pb.Num_UID)) + first := nextUid() + require.Greater(t, first, target) + require.Less(t, first, target+slack) + + // Bumping again to target, or below it, must not move the lease. + require.NoError(t, BumpLease(ctx, zc, target, pb.Num_UID)) + require.NoError(t, BumpLease(ctx, zc, target-500, pb.Num_UID)) + second := nextUid() + require.Greater(t, second, first) + require.Less(t, second, first+slack) +} diff --git a/worker/offline_restore_test.go b/worker/offline_restore_test.go new file mode 100644 index 00000000000..de5f91a8309 --- /dev/null +++ b/worker/offline_restore_test.go @@ -0,0 +1,115 @@ +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package worker + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + "github.com/dgraph-io/badger/v4" + "github.com/dgraph-io/badger/v4/options" + "github.com/dgraph-io/dgraph/v25/posting" + "github.com/dgraph-io/dgraph/v25/protos/pb" + "github.com/dgraph-io/dgraph/v25/x" +) + +// writeTestBackup writes a full backup of group 1 to backupDir using the production +// BackupProcessor. It stores one value posting for every (namespace, uid) pair under the +// predicate "name" and returns the backup's read timestamp. +func writeTestBackup(t *testing.T, backupDir string, nsUids map[uint64][]uint64) uint64 { + t.Helper() + const commitTs, readTs = 5, 10 + + db, err := badger.OpenManaged(badger.DefaultOptions(t.TempDir()).WithLogger(nil)) + require.NoError(t, err) + defer func() { require.NoError(t, db.Close()) }() + + val, err := proto.Marshal(&pb.PostingList{Postings: []*pb.Posting{{ + Uid: 0, Value: []byte("node"), ValType: pb.Posting_STRING, PostingType: pb.Posting_VALUE, + }}}) + require.NoError(t, err) + + var preds []string + txn := db.NewTransactionAt(commitTs, true) + for ns, uids := range nsUids { + attr := x.NamespaceAttr(ns, "name") + preds = append(preds, attr) + for _, uid := range uids { + key := x.DataKey(attr, uid) + require.NoError(t, txn.SetEntry(badger.NewEntry(key, val).WithMeta(posting.BitCompletePosting))) + } + } + require.NoError(t, txn.CommitAt(commitTs, nil)) + + unixTs := time.Now().UTC().Format("20060102.150405.000") + req := &pb.BackupRequest{ + ReadTs: readTs, + GroupId: 1, + UnixTs: unixTs, + Destination: backupDir, + Predicates: preds, + ForceFull: true, + } + bp := NewBackupProcessor(db, req) + defer bp.Close() + _, err = bp.WriteBackup(context.Background()) + require.NoError(t, err) + + m := &Manifest{ + ManifestBase: ManifestBase{ + Type: "full", + ReadTs: readTs, + BackupId: "offline-restore-lease", + BackupNum: 1, + Version: x.ManifestVersion, + Path: "dgraph." + unixTs, + Compression: "snappy", + }, + Groups: map[uint32][]string{1: preds}, + } + require.NoError(t, bp.CompleteBackup(context.Background(), m)) + return readTs +} + +// TestRunOfflineRestoreReportsMaxLease checks that the offline restore (the code path behind +// `dgraph restore`) reports the max UID and namespace present in the backup. `dgraph restore -z` +// uses these values to bump Zero's UID and namespace leases. If they come back as 0, Zero hands +// out UIDs that already belong to restored nodes. +func TestRunOfflineRestoreReportsMaxLease(t *testing.T) { + oldBadger, oldTmp := x.WorkerConfig.Badger, x.WorkerConfig.TmpDir + // BackupProcessor sizes its per-thread state from this; it must match the DB's stream goroutines. + x.WorkerConfig.Badger.NumGoroutines = badger.DefaultOptions("").NumGoroutines + x.WorkerConfig.TmpDir = t.TempDir() + defer func() { x.WorkerConfig.Badger, x.WorkerConfig.TmpDir = oldBadger, oldTmp }() + + const maxUid, maxNs = uint64(0x2711), uint64(2) + nsUids := map[uint64][]uint64{ + x.RootNamespace: {0x1, 0x2, 0x3, maxUid}, + maxNs: {0x10, 0x11}, + } + backupDir := filepath.Join(t.TempDir(), "backups") + readTs := writeTestBackup(t, backupDir, nsUids) + + // Control: the map phase does see the max UID and namespace in the backup. + mapRes, err := RunMapper(&pb.RestoreRequest{Location: backupDir, GroupId: 1, RestoreTs: 1}, t.TempDir()) + require.NoError(t, err) + require.Equal(t, maxUid, mapRes.maxUid, "mapper max uid") + require.Equal(t, maxNs, mapRes.maxNs, "mapper max namespace") + + result := RunOfflineRestore(filepath.Join(t.TempDir(), "p"), backupDir, "", "", nil, options.Snappy, 0) + require.NoError(t, result.Err) + require.Equal(t, readTs, result.Version) + assert.GreaterOrEqual(t, result.MaxLeaseUid, maxUid, + "offline restore must report the max uid in the backup so Zero's uid lease can be bumped") + assert.GreaterOrEqual(t, result.MaxLeaseNsId, maxNs, + "offline restore must report the max namespace in the backup so Zero's ns lease can be bumped") +} diff --git a/worker/online_restore.go b/worker/online_restore.go index a3e9e057989..57aef246e79 100644 --- a/worker/online_restore.go +++ b/worker/online_restore.go @@ -466,23 +466,25 @@ func bumpLease(ctx context.Context, mr *mapResult) error { } zc := pb.NewZeroClient(pl.Get()) - bump := func(val uint64, typ pb.NumLeaseType) error { - _, err := zc.AssignIds(ctx, &pb.Num{Val: val, Type: typ, Bump: true}) - if err != nil && strings.Contains(err.Error(), "Nothing to be leased") { - return nil - } - return err - } - - if err := bump(mr.maxUid, pb.Num_UID); err != nil { + if err := BumpLease(ctx, zc, mr.maxUid, pb.Num_UID); err != nil { return errors.Wrapf(err, "cannot update max uid lease after restore.") } - if err := bump(mr.maxNs, pb.Num_NS_ID); err != nil { + if err := BumpLease(ctx, zc, mr.maxNs, pb.Num_NS_ID); err != nil { return errors.Wrapf(err, "cannot update max namespace lease after restore.") } return nil } +// BumpLease raises Zero's lease of the given type to at least val. It is a no-op when the lease +// already reaches val, so a restore can call it any number of times without consuming ids. +func BumpLease(ctx context.Context, zc pb.ZeroClient, val uint64, typ pb.NumLeaseType) error { + _, err := zc.AssignIds(ctx, &pb.Num{Val: val, Type: typ, Bump: true}) + if err != nil && strings.Contains(err.Error(), "Nothing to be leased") { + return nil + } + return err +} + // create a config object from the request for use with enc package. func getEncConfig(req *pb.RestoreRequest) (*viper.Viper, error) { config := viper.New() @@ -559,6 +561,7 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive, } } + var maxUid, maxNs uint64 for gid := range manifest.Groups { req := &pb.RestoreRequest{ Location: location, @@ -577,9 +580,12 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive, } }() - if _, err := RunMapper(req, mapDir); err != nil { + mapRes, err := RunMapper(req, mapDir) + if err != nil { return LoadResult{Err: errors.Wrap(err, "RunRestore failed to map")} } + maxUid = x.Max(maxUid, mapRes.maxUid) + maxNs = x.Max(maxNs, mapRes.maxNs) pdir := filepath.Join(dir, fmt.Sprintf("p%d", gid)) db, err := badger.OpenManaged(badger.DefaultOptions(pdir). WithCompression(ctype). @@ -609,8 +615,8 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive, return LoadResult{Err: errors.Wrap(err, "RunRestore failed to write group id file")} } } - // TODO: Fix this return value. - return LoadResult{Version: manifest.ValidReadTs()} + // The caller bumps Zero's uid and namespace leases past these, as the online restore does. + return LoadResult{Version: manifest.ValidReadTs(), MaxLeaseUid: maxUid, MaxLeaseNsId: maxNs} } func buildPredsForDefaultNamespace(restorePreds []string, fromNamespace uint64) []string {