Repository: nats-io/nats-server
Pull request #1432 (2020-05-29 17:06:47)

Allow snapshots to optionally check all message checksums
Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/filestore.go b/server/filestore.go
index 64742bfb..b7edb0e8 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -2102,7 +2102,7 @@ func (fs *fileStore) streamSnapshot(w io.WriteCloser, blks []*msgBlock, includeC
 }
 
 // Create a snapshot of this stream and its consumer's state along with messages.
-func (fs *fileStore) Snapshot(deadline time.Duration, includeConsumers bool) (*SnapshotResult, error) {
+func (fs *fileStore) Snapshot(deadline time.Duration, includeConsumers, checkMsgs bool) (*SnapshotResult, error) {
 	fs.mu.Lock()
 	if fs.closed {
 		fs.mu.Unlock()
@@ -2119,6 +2119,13 @@ func (fs *fileStore) Snapshot(deadline time.Duration, includeConsumers bool) (*S
 	blkSize := int(fs.fcfg.BlockSize)
 	fs.mu.Unlock()
 
+	if checkMsgs {
+		bad := fs.checkMsgs()
+		if len(bad) > 0 {
+			return nil, fmt.Errorf("snapshot check detected %d bad messages", len(bad))
+		}
+	}
+
 	pr, pw := net.Pipe()
 
 	// Set a write deadline here to protect ourselves.
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index 0404283b..12a18f9a 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -139,7 +139,7 @@ const (
 	// JSApiRequestNextT is the prefix for the request next message(s) for a consumer in worker/pull mode.
 	JSApiRequestNextT = "$JS.API.CONSUMER.MSG.NEXT.%s.%s"
 
-	// For snapshots and restores
+	// For snapshots and restores. The ack will have additional tokens.
 	jsSnapshotAckT    = "$JS.SNAPSHOT.ACK.%s.%s"
 	jsRestoreDeliverT = "$JS.SNAPSHOT.RESTORE.%s.%s"
 
@@ -314,8 +314,11 @@ type JSApiStreamSnapshotRequest struct {
 	DeliverSubject string `json:"deliver_subject"`
 	// Do not include consumers in the snapshot.
 	NoConsumers bool `json:"no_consumers,omitempty"`
-	// Optional chunk size preference. Otherwise server selects.
+	// Optional chunk size preference.
+	// Best to just let server select.
 	ChunkSize int `json:"chunk_size,omitempty"`
+	// Check all message's checksums prior to snapshot.
+	CheckMsgs bool `json:"jsck,omitempty"`
 }
 
 // JSApiStreamSnapshotResponse is the direct response to the snapshot request.
@@ -1182,19 +1185,23 @@ func (s *Server) jsStreamSnapshotRequest(sub *subscription, c *client, subject,
 		return
 	}
 
-	sr, err := mset.Snapshot(0, !req.NoConsumers)
-	if err != nil {
-		resp.Error = jsError(err)
-		s.sendAPIResponse(c, subject, reply, string(msg), s.jsonResponse(&resp))
-		return
-	}
+	// We will do the snapshot in a go routine as well since check msgs may
+	// stall this go routine.
+	go func() {
+		sr, err := mset.Snapshot(0, !req.NoConsumers, req.CheckMsgs)
+		if err != nil {
+			resp.Error = jsError(err)
+			s.sendAPIResponse(c, subject, reply, string(msg), s.jsonResponse(&resp))
+			return
+		}
 
-	resp.NumBlks = sr.NumBlks
-	resp.BlkSize = sr.BlkSize
-	s.sendAPIResponse(c, subject, reply, string(msg), s.jsonResponse(resp))
+		resp.NumBlks = sr.NumBlks
+		resp.BlkSize = sr.BlkSize
+		s.sendAPIResponse(c, subject, reply, string(msg), s.jsonResponse(resp))
 
-	// Now do the real streaming in a separate go routine.
-	go s.streamSnapshot(c, mset, sr, &req)
+		// Now do the real streaming.
+		s.streamSnapshot(c, mset, sr, &req)
+	}()
 }
 
 // Default chunk size for now.
diff --git a/server/memstore.go b/server/memstore.go
index 69d70d6a..d8da7b77 100644
--- a/server/memstore.go
+++ b/server/memstore.go
@@ -381,7 +381,7 @@ func (ms *memStore) ConsumerStore(_ string, _ *ConsumerConfig) (ConsumerStore, e
 	return &consumerMemStore{ms}, nil
 }
 
-func (ms *memStore) Snapshot(_ time.Duration, _ bool) (*SnapshotResult, error) {
+func (ms *memStore) Snapshot(_ time.Duration, _, _ bool) (*SnapshotResult, error) {
 	return nil, fmt.Errorf("no impl")
 }
 
diff --git a/server/store.go b/server/store.go
index 3af9bb51..1436704b 100644
--- a/server/store.go
+++ b/server/store.go
@@ -63,7 +63,7 @@ type StreamStore interface {
 	Delete() error
 	Stop() error
 	ConsumerStore(name string, cfg *ConsumerConfig) (ConsumerStore, error)
-	Snapshot(deadline time.Duration, includeConsumers bool) (*SnapshotResult, error)
+	Snapshot(deadline time.Duration, includeConsumers, checkMsgs bool) (*SnapshotResult, error)
 }
 
 // RetentionPolicy determines how messages in a set are retained.
diff --git a/server/stream.go b/server/stream.go
index 888f0bd2..3e070a15 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -936,7 +936,7 @@ func (mset *Stream) ackMsg(obs *Consumer, seq uint64) {
 }
 
 // Snapshot creates a snapshot for the stream and possibly consumers.
-func (mset *Stream) Snapshot(deadline time.Duration, includeConsumers bool) (*SnapshotResult, error) {
+func (mset *Stream) Snapshot(deadline time.Duration, checkMsgs, includeConsumers bool) (*SnapshotResult, error) {
 	mset.mu.Lock()
 	if mset.client == nil || mset.store == nil {
 		mset.mu.Unlock()
@@ -954,7 +954,7 @@ func (mset *Stream) Snapshot(deadline time.Duration, includeConsumers bool) (*Sn
 		o.writeState()
 	}
 
-	return store.Snapshot(deadline, includeConsumers)
+	return store.Snapshot(deadline, checkMsgs, includeConsumers)
 }
 
 const snapsDir = "__snapshots__"


Test patch:
diff --git a/server/filestore_test.go b/server/filestore_test.go
index 03002f2d..a2c9ef7e 100644
--- a/server/filestore_test.go
+++ b/server/filestore_test.go
@@ -1116,7 +1116,7 @@ func TestFileStoreSnapshot(t *testing.T) {
 
 	snapshot := func() []byte {
 		t.Helper()
-		r, err := fs.Snapshot(5*time.Second, true)
+		r, err := fs.Snapshot(5*time.Second, true, true)
 		if err != nil {
 			t.Fatalf("Error creating snapshot")
 		}
@@ -1211,8 +1211,7 @@ func TestFileStoreSnapshot(t *testing.T) {
 
 	// Now check to make sure that we get the correct error when trying to delete or erase
 	// a message when a snapshot is in progress and that closing the reader releases that condition.
-
-	sr, err := fs.Snapshot(5*time.Second, true)
+	sr, err := fs.Snapshot(5*time.Second, true, false)
 	if err != nil {
 		t.Fatalf("Error creating snapshot")
 	}
@@ -1233,7 +1232,7 @@ func TestFileStoreSnapshot(t *testing.T) {
 	})
 
 	// Make sure if we do not read properly then it will close the writer and report an error.
-	sr, err = fs.Snapshot(25*time.Millisecond, false)
+	sr, err = fs.Snapshot(25*time.Millisecond, false, false)
 	if err != nil {
 		t.Fatalf("Error creating snapshot")
 	}
diff --git a/test/jetstream_test.go b/test/jetstream_test.go
index f16ecbe5..ce6b94a6 100644
--- a/test/jetstream_test.go
+++ b/test/jetstream_test.go
@@ -2731,7 +2731,7 @@ func TestJetStreamSnapshots(t *testing.T) {
 	// Snapshot state of the stream and consumers.
 	info := info{mset.Config(), mset.State(), obs}
 
-	sr, err := mset.Snapshot(5*time.Second, true)
+	sr, err := mset.Snapshot(5*time.Second, true, false)
 	if err != nil {
 		t.Fatalf("Error getting snapshot: %v", err)
 	}
@@ -2917,7 +2917,6 @@ func TestJetStreamSnapshotsAPI(t *testing.T) {
 	sub, _ := nc.Subscribe(sreq.DeliverSubject, func(m *nats.Msg) {
 		// EOF
 		if len(m.Data) == 0 {
-			m.Sub.Unsubscribe()
 			done <- true
 			return
 		}
@@ -2983,6 +2982,22 @@ func TestJetStreamSnapshotsAPI(t *testing.T) {
 	if mset.State() != state {
 		t.Fatalf("Did not match states, %+v vs %+v", mset.State(), state)
 	}
+
+	// Now ask that the stream be checked first.
+	sreq.ChunkSize = 0
+	sreq.CheckMsgs = true
+	snapshot = snapshot[:0]
+
+	req, _ = json.Marshal(sreq)
+	if _, err = nc.Request(fmt.Sprintf(server.JSApiStreamSnapshotT, mname), req, time.Second); err != nil {
+		t.Fatalf("Unexpected error on snapshot request: %v", err)
+	}
+	// Wait to receive the snapshot.
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("Did not receive our snapshot in time")
+	}
 }
 
 func TestJetStreamSnapshotsAPIPerf(t *testing.T) {