Support containerd in preload endpoint (#324)

* Support prefetch containerd image
uber · Jul 25, 2022 · d4684e2 · d4684e2
1 parent 3a10917
commit d4684e2
Show file tree

Hide file tree

Showing 19 changed files with 890 additions and 106 deletions.
diff --git a/Makefile b/Makefile
@@ -238,7 +238,9 @@ mocks:
 	$(call add_mock,origin/blobclient,ClusterProvider)
 	$(call add_mock,origin/blobclient,ClientResolver)
 
-	$(call add_mock,lib/dockerdaemon,DockerClient)
+	$(call add_mock,lib/containerruntime,Factory)
+	$(call add_mock,lib/containerruntime/containerd,Client)
+	$(call add_mock,lib/containerruntime/dockerdaemon,DockerClient)
 	$(call add_mock,lib/dockerregistry/transfer,ImageTransferer)
 
 	$(call add_mock,tracker/metainfoclient,Client)

diff --git a/agent/agentserver/server.go b/agent/agentserver/server.go
@@ -25,7 +25,7 @@ import (
 
 	"github.com/uber/kraken/build-index/tagclient"
 	"github.com/uber/kraken/core"
-	"github.com/uber/kraken/lib/dockerdaemon"
+	"github.com/uber/kraken/lib/containerruntime"
 	"github.com/uber/kraken/lib/middleware"
 	"github.com/uber/kraken/lib/store"
 	"github.com/uber/kraken/lib/torrent/scheduler"
@@ -41,12 +41,12 @@ type Config struct{}
 
 // Server defines the agent HTTP server.
 type Server struct {
-	config    Config
-	stats     tally.Scope
-	cads      *store.CADownloadStore
-	sched     scheduler.ReloadableScheduler
-	tags      tagclient.Client
-	dockerCli dockerdaemon.DockerClient
+	config           Config
+	stats            tally.Scope
+	cads             *store.CADownloadStore
+	sched            scheduler.ReloadableScheduler
+	tags             tagclient.Client
+	containerRuntime containerruntime.Factory
 }
 
 // New creates a new Server.
@@ -56,13 +56,13 @@ func New(
 	cads *store.CADownloadStore,
 	sched scheduler.ReloadableScheduler,
 	tags tagclient.Client,
-	dockerCli dockerdaemon.DockerClient) *Server {
+	containerRuntime containerruntime.Factory) *Server {
 
 	stats = stats.Tagged(map[string]string{
 		"module": "agentserver",
 	})
 
-	return &Server{config, stats, cads, sched, tags, dockerCli}
+	return &Server{config, stats, cads, sched, tags, containerRuntime}
 }
 
 // Handler returns the HTTP handler.
@@ -166,10 +166,22 @@ func (s *Server) preloadTagHandler(w http.ResponseWriter, r *http.Request) error
 		return handler.Errorf("failed to parse docker image tag")
 	}
 	repo, tag := parts[0], parts[1]
-	if err := s.dockerCli.PullImage(
-		context.Background(), repo, tag); err != nil {
 
-		return handler.Errorf("trigger docker pull: %s", err)
+	rt := httputil.GetQueryArg(r, "runtime", "docker")
+	ns := httputil.GetQueryArg(r, "namespace", "")
+	switch rt {
+	case "docker":
+		if err := s.containerRuntime.DockerClient().
+			PullImage(context.Background(), repo, tag); err != nil {
+			return handler.Errorf("docker pull: %s", err)
+		}
+	case "containerd":
+		if err := s.containerRuntime.ContainerdClient().
+			PullImage(context.Background(), ns, repo, tag); err != nil {
+			return handler.Errorf("containerd pull: %s", err)
+		}
+	default:
+		return handler.Errorf("unsupported container runtime")
 	}
 	return nil
 }

diff --git a/agent/agentserver/server_test.go b/agent/agentserver/server_test.go
@@ -31,7 +31,9 @@ import (
 	"github.com/uber/kraken/lib/torrent/scheduler"
 	"github.com/uber/kraken/lib/torrent/scheduler/connstate"
 	mocktagclient "github.com/uber/kraken/mocks/build-index/tagclient"
-	mockdockerdaemon "github.com/uber/kraken/mocks/lib/dockerdaemon"
+	mockcontainerruntime "github.com/uber/kraken/mocks/lib/containerruntime"
+	mockcontainerd "github.com/uber/kraken/mocks/lib/containerruntime/containerd"
+	mockdockerdaemon "github.com/uber/kraken/mocks/lib/containerruntime/dockerdaemon"
 	mockscheduler "github.com/uber/kraken/mocks/lib/torrent/scheduler"
 	"github.com/uber/kraken/utils/httputil"
 	"github.com/uber/kraken/utils/testutil"
@@ -42,11 +44,13 @@ import (
 )
 
 type serverMocks struct {
-	cads      *store.CADownloadStore
-	sched     *mockscheduler.MockReloadableScheduler
-	tags      *mocktagclient.MockClient
-	dockerCli *mockdockerdaemon.MockDockerClient
-	cleanup   *testutil.Cleanup
+	cads             *store.CADownloadStore
+	sched            *mockscheduler.MockReloadableScheduler
+	tags             *mocktagclient.MockClient
+	dockerCli        *mockdockerdaemon.MockDockerClient
+	containerdCli    *mockcontainerd.MockClient
+	containerRuntime *mockcontainerruntime.MockFactory
+	cleanup          *testutil.Cleanup
 }
 
 func newServerMocks(t *testing.T) (*serverMocks, func()) {
@@ -63,12 +67,15 @@ func newServerMocks(t *testing.T) (*serverMocks, func()) {
 	tags := mocktagclient.NewMockClient(ctrl)
 
 	dockerCli := mockdockerdaemon.NewMockDockerClient(ctrl)
-
-	return &serverMocks{cads, sched, tags, dockerCli, &cleanup}, cleanup.Run
+	containerdCli := mockcontainerd.NewMockClient(ctrl)
+	containerruntime := mockcontainerruntime.NewMockFactory(ctrl)
+	return &serverMocks{
+		cads, sched, tags, dockerCli, containerdCli,
+		containerruntime, &cleanup}, cleanup.Run
 }
 
 func (m *serverMocks) startServer() string {
-	s := New(Config{}, tally.NoopScope, m.cads, m.sched, m.tags, m.dockerCli)
+	s := New(Config{}, tally.NoopScope, m.cads, m.sched, m.tags, m.containerRuntime)
 	addr, stop := testutil.StartServer(s.Handler())
 	m.cleanup.Add(stop)
 	return addr
@@ -262,17 +269,58 @@ func TestDeleteBlobHandler(t *testing.T) {
 }
 
 func TestPreloadHandler(t *testing.T) {
-	require := require.New(t)
-
-	mocks, cleanup := newServerMocks(t)
-	defer cleanup()
+	tag := url.PathEscape("repo1:tag1")
+	tests := []struct {
+		name          string
+		url           string
+		setup         func(*serverMocks)
+		expectedError string
+	}{
+		{
+			name: "success docker",
+			url:  fmt.Sprintf("/preload/tags/%s", tag),
+			setup: func(mocks *serverMocks) {
+				mocks.dockerCli.EXPECT().
+					PullImage(context.Background(), "repo1", "tag1").Return(nil)
+				mocks.containerRuntime.EXPECT().
+					DockerClient().Return(mocks.dockerCli)
+			},
+		},
+		{
+			name: "success containerd",
+			url:  fmt.Sprintf("/preload/tags/%s?runtime=containerd&namespace=name.space1", tag),
+			setup: func(mocks *serverMocks) {
+				mocks.containerdCli.EXPECT().
+					PullImage(context.Background(), "name.space1", "repo1", "tag1").Return(nil)
+				mocks.containerRuntime.EXPECT().
+					ContainerdClient().Return(mocks.containerdCli)
+			},
+		},
+		{
+			name:          "unsupported runtime",
+			url:           fmt.Sprintf("/preload/tags/%s?runtime=crio", tag),
+			setup:         func(_ *serverMocks) {},
+			expectedError: "/preload/tags/repo1:tag1?runtime=crio 500: unsupported container runtime",
+		},
+	}
 
-	addr := mocks.startServer()
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
 
-	tag := url.PathEscape("repo1:tag1")
+			mocks, cleanup := newServerMocks(t)
+			defer cleanup()
 
-	mocks.dockerCli.EXPECT().PullImage(context.Background(), "repo1", "tag1").Return(nil)
+			tt.setup(mocks)
+			addr := mocks.startServer()
 
-	_, err := httputil.Get(fmt.Sprintf("http://%s/preload/tags/%s", addr, tag))
-	require.NoError(err)
+			_, err := httputil.Get(fmt.Sprintf("http://%s%s", addr, tt.url))
+			if tt.expectedError != "" {
+				require.EqualError(err,
+					fmt.Sprintf("GET http://%s%s", addr, tt.expectedError))
+			} else {
+				require.NoError(err)
+			}
+		})
+	}
 }
diff --git a/agent/cmd/cmd.go b/agent/cmd/cmd.go
@@ -22,7 +22,8 @@ import (
 	"github.com/uber/kraken/agent/agentserver"
 	"github.com/uber/kraken/build-index/tagclient"
 	"github.com/uber/kraken/core"
-	"github.com/uber/kraken/lib/dockerdaemon"
+	"github.com/uber/kraken/lib/containerruntime"
+	"github.com/uber/kraken/lib/containerruntime/dockerdaemon"
 	"github.com/uber/kraken/lib/dockerregistry/transfer"
 	"github.com/uber/kraken/lib/store"
 	"github.com/uber/kraken/lib/torrent/networkevent"
@@ -203,13 +204,19 @@ func Run(flags *Flags, opts ...Option) {
 	}
 
 	registryAddr := fmt.Sprintf("127.0.0.1:%d", flags.AgentRegistryPort)
-	dockerCli, err := dockerdaemon.NewDockerClient(config.DockerDaemon, registryAddr)
+	containerRuntimeCfg := config.ContainerRuntime
+	dockerdaemonCfg := dockerdaemon.Config{}
+	if config.DockerDaemon != dockerdaemonCfg {
+		log.Warn("please move docker config under \"container_runtime\"")
+		containerRuntimeCfg.Docker = config.DockerDaemon
+	}
+	containerRuntimeFactory, err := containerruntime.NewFactory(containerRuntimeCfg, registryAddr)
 	if err != nil {
-		log.Fatalf("failed to init docker client for preload: %s", err)
+		log.Fatalf("Failed to create container runtime factory: %s", err)
 	}
 
 	agentServer := agentserver.New(
-		config.AgentServer, stats, cads, sched, tagClient, dockerCli)
+		config.AgentServer, stats, cads, sched, tagClient, containerRuntimeFactory)
 	addr := fmt.Sprintf(":%d", flags.AgentServerPort)
 	log.Infof("Starting agent server on %s", addr)
 	go func() {

diff --git a/agent/cmd/config.go b/agent/cmd/config.go
@@ -16,7 +16,8 @@ package cmd
 import (
 	"github.com/uber/kraken/agent/agentserver"
 	"github.com/uber/kraken/core"
-	"github.com/uber/kraken/lib/dockerdaemon"
+	"github.com/uber/kraken/lib/containerruntime"
+	"github.com/uber/kraken/lib/containerruntime/dockerdaemon"
 	"github.com/uber/kraken/lib/dockerregistry"
 	"github.com/uber/kraken/lib/store"
 	"github.com/uber/kraken/lib/torrent/networkevent"
@@ -31,19 +32,22 @@ import (
 
 // Config defines agent configuration.
 type Config struct {
-	ZapLogging      zap.Config                     `yaml:"zap"`
-	Metrics         metrics.Config                 `yaml:"metrics"`
-	CADownloadStore store.CADownloadStoreConfig    `yaml:"store"`
-	Registry        dockerregistry.Config          `yaml:"registry"`
-	Scheduler       scheduler.Config               `yaml:"scheduler"`
-	PeerIDFactory   core.PeerIDFactory             `yaml:"peer_id_factory"`
-	NetworkEvent    networkevent.Config            `yaml:"network_event"`
-	Tracker         upstream.PassiveHashRingConfig `yaml:"tracker"`
-	BuildIndex      upstream.PassiveConfig         `yaml:"build_index"`
-	AgentServer     agentserver.Config             `yaml:"agentserver"`
-	RegistryBackup  string                         `yaml:"registry_backup"`
-	Nginx           nginx.Config                   `yaml:"nginx"`
-	TLS             httputil.TLSConfig             `yaml:"tls"`
-	AllowedCidrs    []string                       `yaml:"allowed_cidrs"`
-	DockerDaemon    dockerdaemon.Config            `yaml:"docker_daemon"`
+	ZapLogging       zap.Config                     `yaml:"zap"`
+	Metrics          metrics.Config                 `yaml:"metrics"`
+	CADownloadStore  store.CADownloadStoreConfig    `yaml:"store"`
+	Registry         dockerregistry.Config          `yaml:"registry"`
+	Scheduler        scheduler.Config               `yaml:"scheduler"`
+	PeerIDFactory    core.PeerIDFactory             `yaml:"peer_id_factory"`
+	NetworkEvent     networkevent.Config            `yaml:"network_event"`
+	Tracker          upstream.PassiveHashRingConfig `yaml:"tracker"`
+	BuildIndex       upstream.PassiveConfig         `yaml:"build_index"`
+	AgentServer      agentserver.Config             `yaml:"agentserver"`
+	RegistryBackup   string                         `yaml:"registry_backup"`
+	Nginx            nginx.Config                   `yaml:"nginx"`
+	TLS              httputil.TLSConfig             `yaml:"tls"`
+	AllowedCidrs     []string                       `yaml:"allowed_cidrs"`
+	ContainerRuntime containerruntime.Config        `yaml:"container_runtime"`
+
+	// Deprecated
+	DockerDaemon dockerdaemon.Config `yaml:"docker_daemon"`
 }
diff --git a/go.mod b/go.mod
@@ -3,7 +3,8 @@ module github.com/uber/kraken
 go 1.14
 
 require (
-	cloud.google.com/go v0.43.0
+	cloud.google.com/go/storage v1.6.0
+	github.com/Microsoft/hcsshim v0.9.3 // indirect
 	github.com/alecthomas/kingpin v2.2.6+incompatible
 	github.com/alicebob/gopher-json v0.0.0-20180125190556-5a6b3ba71ee6 // indirect
 	github.com/alicebob/miniredis v2.5.0+incompatible
@@ -15,36 +16,37 @@ require (
 	github.com/c2h5oh/datasize v0.0.0-20171227191756-4eba002a5eae
 	github.com/cactus/go-statsd-client v3.1.1+incompatible
 	github.com/cenkalti/backoff v2.2.1+incompatible
-	github.com/docker/distribution v0.0.0-20191024225408-dee21c0394b5
+	github.com/containerd/cgroups v1.0.4 // indirect
+	github.com/containerd/containerd v1.5.7
+	github.com/containerd/continuity v0.0.0-00010101000000-000000000000 // indirect
+	github.com/containerd/fifo v1.0.0 // indirect
+	github.com/docker/distribution v2.7.1+incompatible
 	github.com/docker/docker-credential-helpers v0.6.3
 	github.com/docker/engine-api v0.0.0-20160908232104-4290f40c0566
-	github.com/docker/go-connections v0.0.0-20180821093606-97c2040d34df // indirect
+	github.com/docker/go-events v0.0.0-20190806004212-e31b211e4f1c // indirect
 	github.com/docker/go-metrics v0.0.0-20181218153428-b84716841b82 // indirect
-	github.com/docker/go-units v0.0.0-20181030082039-2fb04c6466a5 // indirect
 	github.com/docker/libtrust v0.0.0-20160708172513-aabc10ec26b7 // indirect
 	github.com/garyburd/redigo v1.6.0
 	github.com/go-chi/chi v4.0.2+incompatible
 	github.com/gofrs/uuid v0.0.0-20190320161447-2593f3d8aa45 // indirect
-	github.com/golang/mock v1.4.4
-	github.com/golang/protobuf v1.3.3
+	github.com/gogo/googleapis v1.4.1 // indirect
+	github.com/golang/mock v1.6.0
+	github.com/golang/protobuf v1.5.0
 	github.com/gomodule/redigo v2.0.0+incompatible // indirect
 	github.com/gorilla/handlers v0.0.0-20190227193432-ac6d24f88de4 // indirect
 	github.com/gorilla/mux v1.7.3
+	github.com/imdario/mergo v0.3.13 // indirect
 	github.com/jackpal/bencode-go v0.0.0-20180813173944-227668e840fa
 	github.com/jinzhu/gorm v1.9.16
 	github.com/jmoiron/sqlx v0.0.0-20190319043955-cdf62fdf55f6
 	github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0 // indirect
-	github.com/konsorten/go-windows-terminal-sequences v1.0.2 // indirect
 	github.com/mattn/go-sqlite3 v1.14.0
-	github.com/opencontainers/go-digest v0.0.0-20190228220655-ac19fd6e7483
+	github.com/opencontainers/go-digest v1.0.0
 	github.com/pressly/goose v2.6.0+incompatible
-	github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90 // indirect
-	github.com/prometheus/common v0.2.0 // indirect
-	github.com/prometheus/procfs v0.0.0-20190328153300-af7bedc223fb // indirect
 	github.com/satori/go.uuid v1.2.0
-	github.com/spaolacci/murmur3 v0.0.0-20170819071325-9f5d223c6079
-	github.com/spf13/cobra v0.0.4 // indirect
-	github.com/stretchr/testify v1.4.0
+	github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72
+	github.com/stretchr/testify v1.7.0
+	github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635 // indirect
 	github.com/uber-go/tally v3.3.11+incompatible
 	github.com/willf/bitset v0.0.0-20190228212526-18bd95f470f9
 	github.com/yuin/gopher-lua v0.0.0-20191128022950-c6266f4fe8d7 // indirect
@@ -53,12 +55,19 @@ require (
 	github.com/yvasiyarov/newrelic_platform_go v0.0.0-20160601141957-9c099fbc30e9 // indirect
 	go.uber.org/atomic v1.5.0
 	go.uber.org/multierr v1.4.0 // indirect
-	go.uber.org/zap v0.0.0-20190327195448-badef736563f
-	golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e
-	golang.org/x/sync v0.0.0-20190423024810-112230192c58
-	golang.org/x/time v0.0.0-20190308202827-9d24e82272b4
-	golang.org/x/tools v0.0.0-20191114200427-caa0b0f7d508 // indirect
-	google.golang.org/api v0.7.0
+	go.uber.org/zap v1.10.0
+	golang.org/x/net v0.0.0-20210825183410-e898025ed96a
+	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
+	golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1
+	google.golang.org/api v0.22.0
 	gopkg.in/validator.v2 v2.0.0-20180514200540-135c24b11c19
-	gopkg.in/yaml.v2 v2.2.2
+	gopkg.in/yaml.v2 v2.3.0
 )
+
+replace github.com/docker/distribution => github.com/docker/distribution v0.0.0-20191024225408-dee21c0394b5
+
+replace github.com/containerd/containerd => github.com/containerd/containerd v1.3.10
+
+replace github.com/containerd/continuity => github.com/containerd/continuity v0.1.0
+
+replace github.com/opencontainers/runc => github.com/opencontainers/runc v1.0.0-rc10