Golang 多协程并发下载 Docker 镜像并实时显示进度

package main

import (
	'context'
	'fmt'
	'io'
	'os'
	'sync'

	'github.com/docker/docker/api/types'
	'github.com/docker/docker/client'
)

func main() {
	// 创建 Docker 客户端
	cli, err := client.NewClientWithOpts(client.FromEnv)
	if err != nil {
		panic(err)
	}

	// 镜像名称和标签
	imageName := 'nginx'
	imageTag := 'latest'

	// 创建一个等待组，用于等待所有协程完成
	var wg sync.WaitGroup

	// 获取镜像的总大小
	imageSize, err := getImageSize(cli, imageName, imageTag)
	if err != nil {
		panic(err)
	}

	// 创建一个通道，用于接收下载进度
	progressCh := make(chan int64)

	// 开启一个协程来显示下载进度
	go func() {
		for {
			downloadedSize := <-progressCh
			fmt.Printf('Downloaded: %d/%d\n', downloadedSize, imageSize)
		}
	}()

	// 开启多个协程进行镜像下载
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			// 下载镜像
			err := pullImage(cli, imageName, imageTag, progressCh)
			if err != nil {
				fmt.Printf('Error downloading image: %s\n', err.Error())
			}
		}()
	}

	// 等待所有协程完成
	wg.Wait()

	// 关闭通道
	close(progressCh)
}

// 获取镜像的总大小
func getImageSize(cli *client.Client, imageName, imageTag string) (int64, error) {
	imageRef := fmt.Sprintf('%s:%s', imageName, imageTag)

	// 查询镜像信息
	imageInspect, _, err := cli.ImageInspectWithRaw(context.Background(), imageRef)
	if err != nil {
		return 0, err
	}

	// 获取镜像大小
	return imageInspect.Size, nil
}

// 下载镜像
func pullImage(cli *client.Client, imageName, imageTag string, progressCh chan<- int64) error {
	imageRef := fmt.Sprintf('%s:%s', imageName, imageTag)

	// 配置镜像下载选项
	options := types.ImagePullOptions{}

	// 发起镜像下载请求
	response, err := cli.ImagePull(context.Background(), imageRef, options)
	if err != nil {
		return err
	}

	// 读取镜像下载流
	defer response.Close()
	_, err = io.Copy(io.MultiWriter(os.Stdout, progressWriter(progressCh)), response)
	if err != nil && err != io.EOF {
		return err
	}

	return nil
}

// 进度写入器
func progressWriter(progressCh chan<- int64) io.Writer {
	var downloadedSize int64

	return &writerWithProgress{
		progressCh:      progressCh,
		downloadedSize:  &downloadedSize,
	}
}

// 带进度的写入器
type writerWithProgress struct {
	progressCh      chan<- int64
	downloadedSize  *int64
}

func (w *writerWithProgress) Write(p []byte) (n int, err error) {
	n = len(p)
	*w.downloadedSize += int64(n)
	w.progressCh <- *w.downloadedSize
	return
}

这段代码展示了如何利用 Golang 的多协程机制以及 Docker SDK 实现高效的 Docker 镜像下载。其关键在于：

1. 多协程并发下载: 通过 sync.WaitGroup 协调多个下载协程，充分利用网络带宽加速下载过程。
2. 实时进度显示: 利用 channel 传递下载进度信息，并在独立协程中实时更新下载进度，提升用户体验。
3. 代码结构清晰: 将镜像大小获取、镜像下载、进度更新等逻辑封装成独立函数，提高代码可读性和可维护性。

通过学习和实践这个例子，你可以更好地掌握 Golang 并发编程和 Docker SDK 的使用，编写出更高效、用户友好的 Docker 镜像下载工具。