godoxy-yusing/src/docker/client.go

package docker

import (
	"net/http"
	"sync"
	"sync/atomic"

	"github.com/docker/cli/cli/connhelper"
	"github.com/docker/docker/client"
	"github.com/sirupsen/logrus"
	"github.com/yusing/go-proxy/common"
	E "github.com/yusing/go-proxy/error"
)

type Client struct {
	key      string
	refCount *atomic.Int32
	*client.Client

	l logrus.FieldLogger
}

func (c Client) DaemonHostname() string {
	url, _ := client.ParseHostURL(c.DaemonHost())
	return url.Hostname()
}

func (c Client) Connected() bool {
	return c.Client != nil
}

// if the client is still referenced, this is no-op
func (c *Client) Close() error {
	if c.refCount.Add(-1) > 0 {
		return nil
	}

	clientMapMu.Lock()
	defer clientMapMu.Unlock()
	delete(clientMap, c.key)

	client := c.Client
	c.Client = nil

	c.l.Debugf("client closed")

	if client != nil {
		return client.Close()
	}
	return nil
}

// ConnectClient creates a new Docker client connection to the specified host.
//
// Returns existing client if available.
//
// Parameters:
//   - host: the host to connect to (either a URL or common.DockerHostFromEnv).
//
// Returns:
//   - Client: the Docker client connection.
//   - error: an error if the connection failed.
func ConnectClient(host string) (Client, E.NestedError) {
	clientMapMu.Lock()
	defer clientMapMu.Unlock()

	// check if client exists
	if client, ok := clientMap[host]; ok {
		client.refCount.Add(1)
		return client, nil
	}

	// create client
	var opt []client.Opt

	switch host {
	case common.DockerHostFromEnv:
		opt = clientOptEnvHost
	default:
		helper, err := E.Check(connhelper.GetConnectionHelper(host))
		if err.HasError() {
			return Client{}, E.UnexpectedError(err.Error())
		}
		if helper != nil {
			httpClient := &http.Client{
				Transport: &http.Transport{
					DialContext: helper.Dialer,
				},
			}
			opt = []client.Opt{
				client.WithHTTPClient(httpClient),
				client.WithHost(helper.Host),
				client.WithAPIVersionNegotiation(),
				client.WithDialContext(helper.Dialer),
			}
		} else {
			opt = []client.Opt{
				client.WithHost(host),
				client.WithAPIVersionNegotiation(),
			}
		}
	}

	client, err := E.Check(client.NewClientWithOpts(opt...))

	if err.HasError() {
		return Client{}, err
	}

	c := Client{
		Client:   client,
		key:      host,
		refCount: &atomic.Int32{},
		l:        logger.WithField("docker_client", client.DaemonHost()),
	}
	c.refCount.Add(1)
	c.l.Debugf("client connected")

	clientMap[host] = c
	return clientMap[host], nil
}

func CloseAllClients() {
	clientMapMu.Lock()
	defer clientMapMu.Unlock()
	for _, client := range clientMap {
		client.Close()
	}
	clientMap = make(map[string]Client)
	logger.Debug("closed all clients")
}

var (
	clientMap        map[string]Client = make(map[string]Client)
	clientMapMu      sync.Mutex
	clientOptEnvHost = []client.Opt{
		client.WithHostFromEnv(),
		client.WithAPIVersionNegotiation(),
	}

	logger = logrus.WithField("module", "docker")
)