anthropic-proxy/internal/embedded/vm.go

package embedded

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/fujin/anthropic-proxy/internal/config"
	"github.com/rs/zerolog/log"
)

type VM struct {
	cfg       config.EmbeddedConfig
	proxyPort int
	cmd       *exec.Cmd
	tmpDir    string
}

func NewVM(cfg config.EmbeddedConfig, proxyPort int) *VM {
	return &VM{cfg: cfg, proxyPort: proxyPort}
}

func (v *VM) Start() error {
	bin, err := ensureBinary("victoria-metrics", v.cfg.VMBinary, v.cfg.BinDir)
	if err != nil {
		return fmt.Errorf("victoria-metrics: %w", err)
	}

	v.tmpDir, err = os.MkdirTemp("", "vm-*")
	if err != nil {
		return fmt.Errorf("create temp dir: %w", err)
	}

	scrapeConfig := fmt.Sprintf(`global:
  scrape_interval: 15s
scrape_configs:
  - job_name: anthropic-proxy
    static_configs:
      - targets:
          - localhost:%d
`, v.proxyPort)

	scrapePath := filepath.Join(v.tmpDir, "scrape.yaml")
	if err := os.WriteFile(scrapePath, []byte(scrapeConfig), 0o644); err != nil {
		return fmt.Errorf("write scrape config: %w", err)
	}

	dataPath := filepath.Join(v.tmpDir, "data")
	if err := os.MkdirAll(dataPath, 0o755); err != nil {
		return fmt.Errorf("create data dir: %w", err)
	}

	v.cmd = exec.Command(bin,
		"-storageDataPath", dataPath,
		"-retentionPeriod", "7d",
		"-httpListenAddr", fmt.Sprintf(":%d", v.cfg.VMPort),
		"-promscrape.config", scrapePath,
	)
	v.cmd.Stdout = &logWriter{level: "info", component: "victoria-metrics"}
	v.cmd.Stderr = &logWriter{level: "error", component: "victoria-metrics"}

	if err := v.cmd.Start(); err != nil {
		return fmt.Errorf("start victoria-metrics: %w", err)
	}

	log.Info().
		Str("binary", bin).
		Int("port", v.cfg.VMPort).
		Int("scrape_target_port", v.proxyPort).
		Msg("victoria-metrics started")

	return nil
}

func (v *VM) Stop() {
	if v.cmd != nil && v.cmd.Process != nil {
		_ = v.cmd.Process.Kill()
		_ = v.cmd.Wait()
	}
	if v.tmpDir != "" {
		_ = os.RemoveAll(v.tmpDir)
	}
}

func (v *VM) Running() bool {
	return v.cmd != nil && v.cmd.Process != nil && v.cmd.ProcessState == nil
}