mesh-catalog/modules/nats/internal/bussnap/restore.go

package bussnap

import (
	"archive/tar"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"nats-tools/internal/busarchive"

	"github.com/nats-io/nats.go"
)

// The keeper's restore in place of a snapshot it took (design 47 G §681): each stream of the file put back under its own
// name, only where no stream of that name is live.

// APIStreamDelete is the subject the keeper removes its own partial stream with.
const APIStreamDelete = "$JS.API.STREAM.DELETE."

const restoreChunk = 64 * 1024

// ManifestOf reads the manifest of a snapshot file: the streams it holds, without unpacking them.
func ManifestOf(path string) (*busarchive.Manifest, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	tr := tar.NewReader(f)
	hdr, err := tr.Next()
	if err != nil {
		return nil, fmt.Errorf("the snapshot does not read: %w", err)
	}
	if hdr.Name != busarchive.ManifestName {
		return nil, fmt.Errorf("the snapshot starts with %q, not its manifest", hdr.Name)
	}
	raw, err := io.ReadAll(io.LimitReader(tr, 64<<20))
	if err != nil {
		return nil, err
	}
	var m busarchive.Manifest
	if err := json.Unmarshal(raw, &m); err != nil {
		return nil, fmt.Errorf("its manifest: %w", err)
	}
	return &m, nil
}

// Live says whether a stream of that name exists.
func Live(ctx context.Context, nc *nats.Conn, name string) (bool, error) {
	err := Request(ctx, nc, APIStreamInfo+name, nil, nil)
	var api *APIError
	switch {
	case err == nil:
		return true, nil
	case errors.As(err, &api) && api.Code == 404:
		return false, nil
	}
	return false, err
}

// DeleteStream removes a stream; one that is not there is no error.
func DeleteStream(ctx context.Context, nc *nats.Conn, name string) error {
	err := Request(ctx, nc, APIStreamDelete+name, nil, nil)
	var api *APIError
	if errors.As(err, &api) && api.Code == 404 {
		return nil
	}
	return err
}

// RestoreFile puts every stream of the snapshot file back under its own name, each sent in chunks the server answers; it
// ends with the context's bound. A stream already there is the caller's to have refused beforehand.
func RestoreFile(ctx context.Context, nc *nats.Conn, path, scratch string, log func(string, ...any)) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	a, err := busarchive.ReadArchive(f, scratch)
	if err != nil {
		return err
	}
	defer a.Close()
	if err := a.Verify(); err != nil {
		return err
	}
	for _, s := range a.Manifest.Streams {
		started := time.Now()
		if err := restoreStream(ctx, nc, a, s); err != nil {
			return fmt.Errorf("stream %s: %w", s.Name, err)
		}
		log("%s: restored in %.2fs", s.Name, time.Since(started).Seconds())
	}
	return nil
}

func restoreStream(ctx context.Context, nc *nats.Conn, a *busarchive.Archive, s busarchive.StreamEntry) error {
	meta, archive, err := a.Open(s)
	if err != nil {
		return err
	}
	defer archive.Close()
	var req map[string]json.RawMessage
	if err := json.Unmarshal(meta, &req); err != nil || req["config"] == nil {
		return errors.New("its meta is not a stream's configuration and state")
	}
	var resp struct {
		DeliverSubject string `json:"deliver_subject"`
	}
	if err := Request(ctx, nc, APIStreamRestore+s.Name, req, &resp); err != nil {
		return err
	}
	if resp.DeliverSubject == "" {
		return errors.New("the server gave nowhere to send the archive")
	}
	buf := make([]byte, restoreChunk)
	for {
		n, err := archive.Read(buf)
		if n > 0 {
			reply, rerr := nc.RequestWithContext(ctx, resp.DeliverSubject, buf[:n])
			if rerr != nil {
				return fmt.Errorf("sending the archive: %w", rerr)
			}
			if len(reply.Data) > 0 {
				return fmt.Errorf("the server refused the archive: %s", Truncate(reply.Data))
			}
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
	}
	final, err := nc.RequestWithContext(ctx, resp.DeliverSubject, nil)
	if err != nil {
		return fmt.Errorf("finishing the restore: %w", err)
	}
	var created struct {
		Error *APIError   `json:"error"`
		State StreamState `json:"state"`
	}
	if err := json.Unmarshal(final.Data, &created); err != nil {
		return fmt.Errorf("the server's last answer is not JSON: %q", Truncate(final.Data))
	}
	if created.Error != nil {
		return created.Error
	}
	if created.State.LastSeq < s.LastSeq || (s.Messages > 0 && created.State.Messages == 0) {
		return fmt.Errorf("restored with %d message(s) up to sequence %d; the snapshot held %d up to %d",
			created.State.Messages, created.State.LastSeq, s.Messages, s.LastSeq)
	}
	return nil
}