magpie/replication.go

package main

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const (
	replicationOpPut    = "put"
	replicationOpDelete = "delete"
)

type ReplicationJob struct {
	ID       int64
	Peer     string
	Op       string
	Key      string
	Attempts int
}

func (s *Server) enqueueReplicationPut(key string) {
	s.enqueueReplication(replicationOpPut, key)
}

func (s *Server) enqueueReplicationDelete(key string) {
	s.enqueueReplication(replicationOpDelete, key)
}

func (s *Server) enqueueReplication(op string, key string) {
	if len(s.cfg.ReplicationPeers) == 0 || len(s.cfg.S3Keys) == 0 || s.cfg.ReplicationSecret == "" {
		return
	}
	for _, peer := range s.cfg.ReplicationPeers {
		if err := s.store.(*FileStore).EnqueueReplication(peer, op, key, s.cfg.ReplicationMaxJobs); err != nil {
			log.Printf("replication enqueue op=%s key=%s peer=%s error=%v", op, sanitizeLogValue(key), sanitizeLogValue(peer), err)
		}
	}
	go s.processReplicationQueue(context.Background())
}

func (s *Server) startReplicationWorker(ctx context.Context) {
	if len(s.cfg.ReplicationPeers) == 0 || len(s.cfg.S3Keys) == 0 || s.cfg.ReplicationSecret == "" {
		return
	}
	go func() {
		ticker := time.NewTicker(2 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				s.processReplicationQueue(ctx)
			}
		}
	}()
}

func (s *Server) processReplicationQueue(ctx context.Context) {
	jobs, err := s.store.(*FileStore).DueReplicationJobs(25)
	if err != nil {
		log.Printf("replication due jobs error=%v", err)
		return
	}
	for _, job := range jobs {
		if ctx.Err() != nil {
			return
		}
		if err := s.performReplicationJob(job); err != nil {
			log.Printf("replication job id=%d op=%s key=%s peer=%s error=%v", job.ID, job.Op, sanitizeLogValue(job.Key), sanitizeLogValue(job.Peer), err)
			if s.cfg.ReplicationMaxAttempts > 0 && job.Attempts+1 >= s.cfg.ReplicationMaxAttempts {
				_ = s.store.(*FileStore).CompleteReplicationJob(job.ID)
				continue
			}
			_ = s.store.(*FileStore).FailReplicationJob(job.ID, job.Attempts+1, err.Error())
			continue
		}
		_ = s.store.(*FileStore).CompleteReplicationJob(job.ID)
	}
}

func (s *Server) performReplicationJob(job ReplicationJob) error {
	switch job.Op {
	case replicationOpPut:
		return s.performReplicationPut(job.Peer, job.Key)
	case replicationOpDelete:
		return s.performReplicationDelete(job.Peer, job.Key)
	default:
		return fmt.Errorf("unknown replication op %q", job.Op)
	}
}

func (s *Server) performReplicationPut(peer string, key string) error {
	obj, err := s.store.Open(key)
	if err != nil {
		return err
	}
	defer obj.Close()

	peerURL, err := replicationObjectURL(peer, key, s.cfg.AllowInsecureReplication)
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodPut, peerURL, obj)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", obj.Info.ContentType)
	req.Header.Set("X-Magpie-Replication", "1")
	if s.cfg.ReplicationSecret != "" {
		req.Header.Set("X-Magpie-Replication-Secret", s.cfg.ReplicationSecret)
	}
	s.signReplication(req)
	resp, err := (&http.Client{Timeout: 5 * time.Minute}).Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("peer status %d", resp.StatusCode)
	}
	return nil
}

func (s *Server) performReplicationDelete(peer string, key string) error {
	peerURL, err := replicationObjectURL(peer, key, s.cfg.AllowInsecureReplication)
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodDelete, peerURL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("X-Magpie-Replication", "1")
	if s.cfg.ReplicationSecret != "" {
		req.Header.Set("X-Magpie-Replication-Secret", s.cfg.ReplicationSecret)
	}
	s.signReplication(req)
	resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("peer status %d", resp.StatusCode)
	}
	return nil
}

func (s *Server) signReplication(req *http.Request) {
	key := s.cfg.S3Keys[0]
	amzDate := time.Now().UTC().Format("20060102T150405Z")
	date := amzDate[:8]
	scope := date + "/" + s.cfg.S3Region + "/s3/aws4_request"
	signedHeaders := "host;x-amz-content-sha256;x-amz-date;x-magpie-replication"
	if req.Header.Get("X-Magpie-Replication-Secret") != "" {
		signedHeaders += ";x-magpie-replication-secret"
	}
	req.Header.Set("X-Amz-Date", amzDate)
	req.Header.Set("X-Amz-Content-Sha256", "UNSIGNED-PAYLOAD")
	canonical := canonicalRequest(req, signedHeaders, "UNSIGNED-PAYLOAD", false)
	sig := signature(key.Secret, amzDate, scope, canonical)
	req.Header.Set("Authorization", awsAlgorithm+" Credential="+key.ID+"/"+scope+", SignedHeaders="+signedHeaders+", Signature="+sig)
}

func replicationObjectURL(peer string, key string, allowInsecure bool) (string, error) {
	if err := ValidateKey(key); err != nil {
		return "", err
	}
	if err := validateReplicationPeer(peer, allowInsecure); err != nil {
		return "", err
	}
	parsed, err := url.Parse(peer)
	if err != nil {
		return "", err
	}
	parsed.RawQuery = ""
	parsed.Fragment = ""
	parsed.Path = strings.TrimRight(parsed.Path, "/") + "/" + key
	parsed.RawPath = ""
	return parsed.String(), nil
}

func validateReplicationPeer(peer string, allowInsecure bool) error {
	parsed, err := url.Parse(peer)
	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
		return fmt.Errorf("invalid replication peer %q", peer)
	}
	if parsed.RawQuery != "" || parsed.Fragment != "" {
		return fmt.Errorf("replication peer %q must not include query or fragment", peer)
	}
	if parsed.Scheme == "https" {
		return nil
	}
	if allowInsecure && parsed.Scheme == "http" {
		return nil
	}
	return fmt.Errorf("replication peer %q must use https or set MAGPIE_ALLOW_INSECURE_REPLICATION=true", peer)
}

func (s *FileStore) EnqueueReplication(peer string, op string, key string, maxJobs int) error {
	now := time.Now().UTC().Format(time.RFC3339Nano)
	result, err := s.db.Exec(`
UPDATE replication_jobs
SET next_attempt_at = ?, updated_at = ?
WHERE peer = ? AND op = ? AND key = ?
`, now, now, peer, op, key)
	if err != nil {
		return err
	}
	if rows, err := result.RowsAffected(); err == nil && rows > 0 {
		return nil
	}
	if maxJobs <= 0 {
		maxJobs = defaultReplicationMaxJobs
	}
	var pending int
	if err := s.db.QueryRow(`SELECT COUNT(*) FROM replication_jobs`).Scan(&pending); err != nil {
		return err
	}
	if pending >= maxJobs {
		return fmt.Errorf("replication queue is full")
	}
	_, err = s.db.Exec(`
INSERT INTO replication_jobs (peer, op, key, next_attempt_at, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?)
`, peer, op, key, now, now, now)
	return err
}

func (s *FileStore) DueReplicationJobs(limit int) ([]ReplicationJob, error) {
	if limit <= 0 || limit > 100 {
		limit = 25
	}
	now := time.Now().UTC().Format(time.RFC3339Nano)
	rows, err := s.db.Query(`
SELECT id, peer, op, key, attempts
FROM replication_jobs
WHERE next_attempt_at <= ?
ORDER BY id ASC
LIMIT ?
`, now, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	jobs := make([]ReplicationJob, 0)
	for rows.Next() {
		var job ReplicationJob
		if err := rows.Scan(&job.ID, &job.Peer, &job.Op, &job.Key, &job.Attempts); err != nil {
			return nil, err
		}
		jobs = append(jobs, job)
	}
	return jobs, rows.Err()
}

func (s *FileStore) CompleteReplicationJob(id int64) error {
	_, err := s.db.Exec(`DELETE FROM replication_jobs WHERE id = ?`, id)
	return err
}

func (s *FileStore) FailReplicationJob(id int64, attempts int, message string) error {
	delay := time.Duration(attempts*attempts) * time.Second
	if delay > 5*time.Minute {
		delay = 5 * time.Minute
	}
	next := time.Now().UTC().Add(delay).Format(time.RFC3339Nano)
	now := time.Now().UTC().Format(time.RFC3339Nano)
	_, err := s.db.Exec(`
UPDATE replication_jobs
SET attempts = ?, next_attempt_at = ?, last_error = ?, updated_at = ?
WHERE id = ?
`, attempts, next, truncateError(message), now, id)
	return err
}

func (s *FileStore) EnqueueRepair(peers []string) (int, error) {
	if len(peers) == 0 {
		return 0, nil
	}
	count := 0
	after := ""
	for {
		page, err := s.ListPage("", after, 1000)
		if err != nil {
			return count, err
		}
		for _, object := range page.Objects {
			for _, peer := range peers {
				if err := s.EnqueueReplication(peer, replicationOpPut, object.Key, defaultReplicationMaxJobs); err != nil {
					return count, err
				}
				count++
			}
		}
		if page.NextAfter == "" {
			return count, nil
		}
		after = page.NextAfter
	}
}

func (s *FileStore) ReplicationQueueStats() (ReplicationQueueStats, error) {
	var stats ReplicationQueueStats
	err := s.db.QueryRow(`SELECT COUNT(*), COALESCE(MAX(attempts), 0) FROM replication_jobs`).Scan(&stats.PendingJobs, &stats.MaxAttempts)
	if errors.Is(err, sql.ErrNoRows) {
		return stats, nil
	}
	return stats, err
}

func truncateError(message string) string {
	if len(message) <= 500 {
		return message
	}
	return message[:500]
}

type ReplicationQueueStats struct {
	PendingJobs int64 `json:"pending_jobs"`
	MaxAttempts int   `json:"max_attempts"`
}