qq_bot/service/exec/exec.go

package exec

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"time"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/image"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
)

var DockerContainer *dockerContainer

func init() {
	//检测docker服务是否已运行
	go func() {
		maxRetries := 5
		retryCount := 0
		for {
			// Initialize DockerContainer with memory, CPU, and disk limits
			container, err := NewDockerContainer(int64(2*1024*1024*1024), int64(1000000000))
			if err != nil {
				retryCount++
				if retryCount > maxRetries {
					log.Fatalf("Failed to initialize Docker container after %d retries: %v", maxRetries, err)
				}
				log.Printf("NewDockerContainer failed (attempt %d/%d): %v", retryCount, maxRetries, err)
				time.Sleep(2 * time.Second)
				continue
			}

			if err := container.CreateAndStartContainer(); err != nil {
				retryCount++
				if retryCount > maxRetries {
					log.Printf("Failed to create and start container after %d retries: %v", maxRetries, err)
					return
				}
				log.Printf("CreateAndStartContainer failed (attempt %d/%d): %v", retryCount, maxRetries, err)
				time.Sleep(2 * time.Second)
				continue
			}

			DockerContainer = container
			log.Println("Docker container successfully initialized")
			return
		}
	}()
}

type dockerContainer struct {
	Client      *client.Client
	Ctx         context.Context
	ContainerID string
	Memory      int64 // Memory in bytes
	CPU         int64 // CPU in NanoCPUs
}

// NewDockerContainer initializes Docker client, context, and resource limits
func NewDockerContainer(memory int64, cpu int64) (*dockerContainer, error) {
	ctx := context.Background()
	cli, err := client.NewClientWithOpts(client.FromEnv)
	if err != nil {
		return nil, err
	}
	cli.NegotiateAPIVersion(ctx)

	return &dockerContainer{
		Client: cli,
		Ctx:    ctx,
		Memory: memory,
		CPU:    cpu,
	}, nil
}

// CreateAndStartContainer creates a container, sets resource limits, and starts it
// The container runs `tail -f /dev/null` to keep running, waiting for new commands
func (dc *dockerContainer) CreateAndStartContainer() error {
	if dc == nil {
		return fmt.Errorf("no container found, please create and start the container first")
	}
	// 如果容器已经存在，则跳过创建
	if dc.ContainerID != "" {
		log.Printf("Container already exists, skipping creation.")
		return nil
	}

	// Pull debian image if not available locally
	_, err := dc.Client.ImagePull(dc.Ctx, "python:3.13", image.PullOptions{})
	if err != nil {
		return err
	}

	// Define resource limits based on DockerContainer struct fields
	hostConfig := &container.HostConfig{
		Resources: container.Resources{
			Memory:   dc.Memory, // Set memory limit
			NanoCPUs: dc.CPU,    // Set CPU limit
		},
	}

	// Use the custom network to isolate the container
	networkConfig := &network.NetworkingConfig{
		EndpointsConfig: map[string]*network.EndpointSettings{
			"qq_bot_net": {}, // Attach container to the custom network
		},
	}

	// Create a container that runs `tail -f /dev/null` to keep it alive
	resp, err := dc.Client.ContainerCreate(dc.Ctx, &container.Config{
		Image: "python:3.13",
		Cmd:   []string{"tail", "-f", "/dev/null"}, // Keep container running
		Tty:   true,
	}, hostConfig, networkConfig, nil, "")
	if err != nil {
		return err
	}

	// Store container ID in the struct
	dc.ContainerID = resp.ID

	// Start the container
	if err := dc.Client.ContainerStart(dc.Ctx, resp.ID, container.StartOptions{}); err != nil {
		return err
	}

	log.Printf("Container created and started, waiting for commands...")
	return nil
}

// ExecCommandInContainer executes a command in the existing container and returns its output
func (dc *dockerContainer) ExecCommandInContainer(command string) (string, error) {
	if dc == nil {
		return "", fmt.Errorf("no container found, please create and start the container first")
	}

	execConfig := container.ExecOptions{
		Cmd:          []string{"sh", "-c", command},
		AttachStdout: true,
		AttachStderr: true,
		Tty:          true,
	}
	var cancel context.CancelFunc
	timeout := 10 * time.Second
	ctxWithTimeout, cancel := context.WithTimeout(dc.Ctx, timeout)
	defer cancel()

	// 创建一个 exec 实例
	execIDResp, err := dc.Client.ContainerExecCreate(ctxWithTimeout, dc.ContainerID, execConfig)
	if err != nil {
		return "", err
	}

	// 开始执行命令
	resp, err := dc.Client.ContainerExecAttach(ctxWithTimeout, execIDResp.ID, container.ExecStartOptions{
		Tty: true,
	})
	if err != nil {
		return "", err
	}
	defer resp.Close()

	// Capture the output
	var outputBuffer bytes.Buffer
	_, err = io.Copy(&outputBuffer, resp.Reader)
	if err != nil {
		return "", err
	}

	return outputBuffer.String(), nil
}

// RestartAndCleanContainer stops, removes, and recreates the container
func (dc *dockerContainer) RestartAndCleanContainer() error {
	if dc == nil {
		return nil
	}
	if dc.ContainerID == "" {
		return fmt.Errorf("no container to restart")
	}

	// Stop the container
	timeout := 10 * time.Second
	timeoutSecs := int(timeout.Seconds())
	if err := dc.Client.ContainerStop(dc.Ctx, dc.ContainerID, container.StopOptions{Timeout: &timeoutSecs}); err != nil {
		return err
	}

	// Remove the container to clean it
	if err := dc.Client.ContainerRemove(dc.Ctx, dc.ContainerID, container.RemoveOptions{Force: true}); err != nil {
		return err
	}

	// Recreate the container with the same resource limits
	resp, err := dc.Client.ContainerCreate(dc.Ctx, &container.Config{
		Image: "python:3.13",
		Cmd:   []string{"tail", "-f", "/dev/null"},
		Tty:   true,
	}, &container.HostConfig{
		Resources: container.Resources{
			Memory:   dc.Memory, // Reuse memory limit
			NanoCPUs: dc.CPU,    // Reuse CPU limit
		},
	}, nil, nil, "")
	if err != nil {
		return err
	}

	// Store new container ID
	dc.ContainerID = resp.ID

	// Start the new container
	if err := dc.Client.ContainerStart(dc.Ctx, resp.ID, container.StartOptions{}); err != nil {
		return err
	}

	log.Printf("Container restarted and cleaned.")
	return nil
}

// RemoveContainer stops and removes the container
func (dc *dockerContainer) RemoveContainer() error {
	if dc == nil {
		return nil
	}
	if dc.ContainerID == "" {
		return fmt.Errorf("no container to remove")
	}

	// Stop the container
	timeout := 10 * time.Second
	timeoutSecs := int(timeout.Seconds())
	if err := dc.Client.ContainerStop(dc.Ctx, dc.ContainerID, container.StopOptions{Timeout: &timeoutSecs}); err != nil {
		return err
	}

	// Remove the container
	if err := dc.Client.ContainerRemove(dc.Ctx, dc.ContainerID, container.RemoveOptions{Force: true}); err != nil {
		return err
	}

	log.Printf("Container removed.")
	return nil
}