next.orly.dev/cmd/benchmark/dgraph_docker.go

package main

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

// DgraphDocker manages a dgraph instance via Docker Compose
type DgraphDocker struct {
	composeFile string
	projectName string
	running     bool
}

// NewDgraphDocker creates a new dgraph Docker manager
func NewDgraphDocker() *DgraphDocker {
	// Try to find the docker-compose file in the current directory first
	composeFile := "docker-compose-dgraph.yml"

	// If not found, try the cmd/benchmark directory (for running from project root)
	if _, err := os.Stat(composeFile); os.IsNotExist(err) {
		composeFile = filepath.Join("cmd", "benchmark", "docker-compose-dgraph.yml")
	}

	return &DgraphDocker{
		composeFile: composeFile,
		projectName: "orly-benchmark-dgraph",
		running:     false,
	}
}

// Start starts the dgraph Docker containers
func (d *DgraphDocker) Start(ctx context.Context) error {
	fmt.Println("Starting dgraph Docker containers...")

	// Stop any existing containers first
	d.Stop()

	// Start containers
	cmd := exec.CommandContext(
		ctx,
		"docker-compose",
		"-f", d.composeFile,
		"-p", d.projectName,
		"up", "-d",
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to start dgraph containers: %w", err)
	}

	fmt.Println("Waiting for dgraph to be healthy...")

	// Wait for health checks to pass
	if err := d.waitForHealthy(ctx, 60*time.Second); err != nil {
		d.Stop() // Clean up on failure
		return err
	}

	d.running = true
	fmt.Println("Dgraph is ready!")
	return nil
}

// waitForHealthy waits for dgraph to become healthy
func (d *DgraphDocker) waitForHealthy(ctx context.Context, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)

	for time.Now().Before(deadline) {
		// Check if alpha is healthy by checking docker health status
		cmd := exec.CommandContext(
			ctx,
			"docker",
			"inspect",
			"--format={{.State.Health.Status}}",
			"orly-benchmark-dgraph-alpha",
		)

		output, err := cmd.Output()
		if err == nil && string(output) == "healthy\n" {
			// Additional short wait to ensure full readiness
			time.Sleep(2 * time.Second)
			return nil
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
			// Continue waiting
		}
	}

	return fmt.Errorf("dgraph failed to become healthy within %v", timeout)
}

// Stop stops and removes the dgraph Docker containers
func (d *DgraphDocker) Stop() error {
	if !d.running {
		// Try to stop anyway in case of untracked state
		cmd := exec.Command(
			"docker-compose",
			"-f", d.composeFile,
			"-p", d.projectName,
			"down", "-v",
		)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		_ = cmd.Run() // Ignore errors
		return nil
	}

	fmt.Println("Stopping dgraph Docker containers...")

	cmd := exec.Command(
		"docker-compose",
		"-f", d.composeFile,
		"-p", d.projectName,
		"down", "-v",
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to stop dgraph containers: %w", err)
	}

	d.running = false
	fmt.Println("Dgraph containers stopped")
	return nil
}

// GetGRPCEndpoint returns the dgraph gRPC endpoint
func (d *DgraphDocker) GetGRPCEndpoint() string {
	return "localhost:9080"
}

// IsRunning returns whether dgraph is running
func (d *DgraphDocker) IsRunning() bool {
	return d.running
}

// Logs returns the logs from dgraph containers
func (d *DgraphDocker) Logs() error {
	cmd := exec.Command(
		"docker-compose",
		"-f", d.composeFile,
		"-p", d.projectName,
		"logs",
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}