Merge release-20190806.1-358-ga99d347 (automated)

4 files changed, 24 insertions, 2 deletions

diff --git a/pkg/state/decode.go b/pkg/state/decode.go
index 47e6b878a..590c241a3 100644
--- a/pkg/state/decode.go
+++ b/pkg/state/decode.go
@@ -16,6 +16,7 @@ package state
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -133,6 +134,9 @@ func (os *objectState) findCycle() []*objectState {
 // to ensure that all callbacks are executed, otherwise the callback graph was
 // not acyclic.
 type decodeState struct {
+	// ctx is the decode context.
+	ctx context.Context
+
 	// objectByID is the set of objects in progress.
 	objectsByID map[uint64]*objectState
 
diff --git a/pkg/state/encode.go b/pkg/state/encode.go
index 5d9409a45..c5118d3a9 100644
--- a/pkg/state/encode.go
+++ b/pkg/state/encode.go
@@ -16,6 +16,7 @@ package state
 
 import (
 	"container/list"
+	"context"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -38,6 +39,9 @@ type queuedObject struct {
 // The encoding process is a breadth-first traversal of the object graph. The
 // inherent races and dependencies are much simpler than the decode case.
 type encodeState struct {
+	// ctx is the encode context.
+	ctx context.Context
+
 	// lastID is the last object ID.
 	//
 	// See idsByObject for context. Because of the special zero encoding
diff --git a/pkg/state/map.go b/pkg/state/map.go
index 7e6fefed4..4f3ebb0da 100644
--- a/pkg/state/map.go
+++ b/pkg/state/map.go
@@ -15,6 +15,7 @@
 package state
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"sort"
@@ -219,3 +220,13 @@ func (m Map) AfterLoad(fn func()) {
 	// data dependencies have been cleared.
 	m.os.callbacks = append(m.os.callbacks, fn)
 }
+
+// Context returns the current context object.
+func (m Map) Context() context.Context {
+	if m.es != nil {
+		return m.es.ctx
+	} else if m.ds != nil {
+		return m.ds.ctx
+	}
+	return context.Background() // No context.
+}
diff --git a/pkg/state/state.go b/pkg/state/state.go
index d408ff84a..dbe507ab4 100644
--- a/pkg/state/state.go
+++ b/pkg/state/state.go
@@ -50,6 +50,7 @@
 package state
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"reflect"
@@ -86,9 +87,10 @@ func UnwrapErrState(err error) error {
 }
 
 // Save saves the given object state.
-func Save(w io.Writer, rootPtr interface{}, stats *Stats) error {
+func Save(ctx context.Context, w io.Writer, rootPtr interface{}, stats *Stats) error {
 	// Create the encoding state.
 	es := &encodeState{
+		ctx:         ctx,
 		idsByObject: make(map[uintptr]uint64),
 		w:           w,
 		stats:       stats,
@@ -101,9 +103,10 @@ func Save(w io.Writer, rootPtr interface{}, stats *Stats) error {
 }
 
 // Load loads a checkpoint.
-func Load(r io.Reader, rootPtr interface{}, stats *Stats) error {
+func Load(ctx context.Context, r io.Reader, rootPtr interface{}, stats *Stats) error {
 	// Create the decoding state.
 	ds := &decodeState{
+		ctx:         ctx,
 		objectsByID: make(map[uint64]*objectState),
 		deferred:    make(map[uint64]*pb.Object),
 		r:           r,