mesh-catalog/modules/nats/cmd/mesh-nats-snapshot/restore.go

package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/nats-io/nats.go"
	"nats-tools/internal/busarchive"
)

// restoreChunk is what is sent at a time; the server answers each before the next goes.
const restoreChunk = 64 * 1024

// RestoreAll puts every stream of the archive into the server nc reaches, one at a time, and holds
// each restored stream's state to the manifest. **Nothing is restored over anything**: a server that
// already holds any of the streams is refused before the first is sent.
func RestoreAll(ctx context.Context, nc *nats.Conn, a *busarchive.Archive, log func(string, ...any)) error {
	var present []string
	for _, s := range a.Manifest.Streams {
		err := request(ctx, nc, apiStreamInfo+s.Name, nil, nil)
		var api *apiError
		switch {
		case err == nil:
			present = append(present, s.Name)
		case errors.As(err, &api) && api.Code == 404:
		default:
			return err
		}
	}
	if len(present) > 0 {
		return fmt.Errorf("the server already holds %s; nothing is restored over a live stream — restore into a new "+
			"store (--into) and swap it in with the bus stopped", strings.Join(present, ", "))
	}
	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, at least %d message(s) to sequence %d, in %.2fs", s.Name, s.Messages, s.LastSeq, 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 {
			cctx, cancel := context.WithTimeout(ctx, 30*time.Second)
			reply, rerr := nc.RequestWithContext(cctx, resp.DeliverSubject, buf[:n])
			cancel()
			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
		}
	}
	// The end: an empty message, answered once the server has rebuilt the stream from what it was
	// sent — which for a large stream takes a while.
	fctx, cancel := context.WithTimeout(ctx, 30*time.Minute)
	defer cancel()
	final, err := nc.RequestWithContext(fctx, 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
	}
	return heldTo(created.State, s)
}

// heldTo is whether a restored stream holds at least what its snapshot said it held.
//
// **At least, not exactly — found by the restore test, not assumed.** The state the server answers a
// snapshot request with is taken when the snapshot starts; the stream goes on taking messages while
// its blocks are read out, and a block read a moment later carries some of what was appended in that
// moment. So a stream written during its snapshot restores with a later last sequence than the
// manifest says (90075 where it said 90024, with a writer publishing every 2ms through it), and of
// the messages in that tail some may be there and some not. Every message up to the manifest's last
// sequence is there, exactly; that is the snapshot's promise, and what the test holds it to message
// by message. What must never happen is a restored stream that ends before the manifest said it did.
func heldTo(restored streamState, s busarchive.StreamEntry) error {
	if restored.LastSeq < s.LastSeq || (s.Messages > 0 && restored.Messages == 0) {
		return fmt.Errorf("restored with %d message(s) up to sequence %d; the snapshot held %d up to %d",
			restored.Messages, restored.LastSeq, s.Messages, s.LastSeq)
	}
	return nil
}

// RestoreOffline fills a new store directory with the archive's streams, by a server of this
// program's own: nats-server, the same binary the bus runs, on loopback, with the bus's one account
// and a user that exists only for this run. It is stopped when the streams are in. What it leaves
// is a store directory laid out exactly as the bus's — `jetstream/<account>/streams/<name>` — which
// a person swaps in for the bus's own with the bus stopped. It answers that jetstream directory.
//
// Why a server of its own rather than the live bus: a stream is restored only where it does not
// exist, and on the live bus every stream exists — the controller asserts them on every start. So a
// restore beside the live data, swapped in by a person, is the only one that never touches it.
func RestoreOffline(ctx context.Context, a *busarchive.Archive, into, account string, log func(string, ...any)) (string, error) {
	if !busarchive.ValidName(account) {
		return "", fmt.Errorf("%q is not an account name", account)
	}
	if entries, err := os.ReadDir(into); err == nil && len(entries) > 0 {
		return "", fmt.Errorf("%s is not empty; a restore goes into a new directory, never over one", into)
	}
	if err := os.MkdirAll(into, 0o700); err != nil {
		return "", err
	}
	store, err := filepath.Abs(into)
	if err != nil {
		return "", err
	}
	binary, err := exec.LookPath("nats-server")
	if err != nil {
		return "", errors.New("no nats-server here to restore with; run this in the bus's own image")
	}
	work, err := os.MkdirTemp("", "mesh-nats-restore-server-*")
	if err != nil {
		return "", err
	}
	defer os.RemoveAll(work)
	port, err := freePort()
	if err != nil {
		return "", err
	}
	password := randomHex(24)
	conf := filepath.Join(work, "restore.conf")
	config := fmt.Sprintf(`listen: "127.0.0.1:%d"
jetstream { store_dir: %q }
accounts { %s { jetstream: enabled, users: [ { user: "restore", password: %q } ] } }
`, port, store, account, password)
	if err := os.WriteFile(conf, []byte(config), 0o600); err != nil {
		return "", err
	}
	cmd := exec.Command(binary, "-c", conf)
	logFile, err := os.Create(filepath.Join(work, "server.log"))
	if err != nil {
		return "", err
	}
	defer logFile.Close()
	cmd.Stdout, cmd.Stderr = logFile, logFile
	if err := cmd.Start(); err != nil {
		return "", err
	}
	exited := make(chan error, 1)
	go func() { exited <- cmd.Wait() }()
	stop := func() error {
		_ = cmd.Process.Signal(syscall.SIGTERM)
		select {
		case err := <-exited:
			return err
		case <-time.After(30 * time.Second):
			_ = cmd.Process.Kill()
			<-exited
			return errors.New("the restore's server did not stop within 30s and was killed")
		}
	}
	serverLog := func() string {
		raw, _ := os.ReadFile(filepath.Join(work, "server.log"))
		return strings.TrimSpace(string(raw))
	}

	url := fmt.Sprintf("nats://127.0.0.1:%d", port)
	var nc *nats.Conn
	for deadline := time.Now().Add(20 * time.Second); ; {
		nc, err = connect(url, Credential{User: "restore", Password: password})
		if err == nil {
			break
		}
		if time.Now().After(deadline) {
			_ = stop()
			return "", fmt.Errorf("the restore's server did not come up: %v\n%s", err, serverLog())
		}
		select {
		case err := <-exited:
			return "", fmt.Errorf("the restore's server stopped: %v\n%s", err, serverLog())
		case <-time.After(200 * time.Millisecond):
		}
	}
	err = RestoreAll(ctx, nc, a, log)
	nc.Close()
	if stopErr := stop(); err == nil && stopErr != nil && !isTerminated(stopErr) {
		err = stopErr
	}
	if err != nil {
		return "", err
	}
	return filepath.Join(store, "jetstream"), nil
}

// isTerminated is a server that stopped because it was asked to.
func isTerminated(err error) bool {
	var exit *exec.ExitError
	if errors.As(err, &exit) {
		if status, ok := exit.Sys().(syscall.WaitStatus); ok && status.Signaled() && status.Signal() == syscall.SIGTERM {
			return true
		}
	}
	return false
}

func freePort() (int, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}

func randomHex(n int) string {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return hex.EncodeToString(b)
}