From dcddca6e132b3417d01e3cf57a32d091678b4fcb Mon Sep 17 00:00:00 2001 From: Lloyd Mangnall Date: Wed, 21 Oct 2026 14:57:17 -0600 Subject: [PATCH 1/2] fix(restore): report max uid and namespace from offline restore RunOfflineRestore discarded the map phase's result, so its LoadResult always carried MaxLeaseUid and MaxLeaseNsId of 0. `dgraph restore -z` treats 0 as "backup of an empty DB" and skips the Zero lease bumps, so a cluster restored offline from a non-empty backup keeps maxUID at 0 and hands out uids that already belong to restored nodes: the first new node can land on 0x1 and merge into an existing one. The namespace lease is skipped the same way. Keep the mapper's maxUid and maxNs across groups and return them, as the online restore already does when it bumps the leases. Co-Authored-By: Claude Opus 5.5 --- worker/offline_restore_test.go | 115 +++++++++++++++++++++++++++++++++ worker/online_restore.go | 10 ++- 2 files changed, 122 insertions(+), 3 deletions(-) create mode 100644 worker/offline_restore_test.go 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..07d73f3690f 100644 --- a/worker/online_restore.go +++ b/worker/online_restore.go @@ -559,6 +559,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 +578,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 +613,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 { From 0875c182faac3fd5d310c876f6a2faf399fbf8c9 Mon Sep 17 00:00:00 2001 From: Lloyd Mangnall Date: Sat, 24 Oct 2026 22:52:12 -0600 Subject: [PATCH 2/2] fix(restore): bump leases after offline restore instead of leasing more dgraph restore -z sent the restored max uid and namespace to Zero as a plain AssignIds request, which leases that many new ids rather than raising the lease to that value. On a fresh Zero the result is the same, but every further restore against the same Zero moved the lease again, by the backup's max each time. Move the online restore's bump into worker.BumpLease, which raises the lease to at least the given value and treats "Nothing to be leased" as a no-op, and use it from both restore paths. TestBumpLeaseIsIdempotent (integration) checks that a repeated bump, or one below the current lease, leaves it where it is. Co-Authored-By: Claude Opus 5.5 --- backup/run.go | 4 ++- worker/lease_test.go | 59 ++++++++++++++++++++++++++++++++++++++++ worker/online_restore.go | 22 ++++++++------- 3 files changed, 74 insertions(+), 11 deletions(-) create mode 100644 worker/lease_test.go 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/online_restore.go b/worker/online_restore.go index 07d73f3690f..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()