feat: operator scalable

Author: morlay

pipeline/operator/container/image_registry.go

@@ -0,0 +1,45 @@
+package container
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"net/url"
+
+	"github.com/docker/docker/api/types"
+)
+
+func ParseImageRegistry(uri string) (*ImageRegistry, error) {
+	u, err := url.ParseRequestURI(uri)
+	if err != nil {
+		return nil, err
+	}
+
+	endpoint := &ImageRegistry{}
+
+	endpoint.Host = u.Host
+	endpoint.Prefix = u.Path
+
+	if u.User != nil {
+		endpoint.Username = u.User.Username()
+		endpoint.Password, _ = u.User.Password()
+	}
+
+	return endpoint, nil
+}
+
+type ImageRegistry struct {
+	Host     string
+	Username string
+	Password string
+	Prefix   string
+}
+
+func (s ImageRegistry) Fix(imageRef string) string {
+	return s.Host + s.Prefix + imageRef
+}
+
+func (s ImageRegistry) RegistryAuth() string {
+	authConfig := types.AuthConfig{Username: s.Username, Password: s.Password, ServerAddress: s.Host}
+	b, _ := json.Marshal(authConfig)
+	return base64.StdEncoding.EncodeToString(b)
+}

pipeline/operator/container/operatormgr.go

@@ -0,0 +1,39 @@
+package container
+
+import (
+	"context"
+
+	"github.com/querycap/pipeline/pipeline"
+	"github.com/querycap/pipeline/spec"
+)
+
+func NewOperatorMgr(podController PodController, container *Container) pipeline.OperatorMgr {
+	return &operatorMgr{
+		podController: podController,
+		container:     container,
+	}
+}
+
+type operatorMgr struct {
+	podController PodController
+	container     *Container
+}
+
+func (d *operatorMgr) Up(scope string, stage string, step spec.Stage, replicas int) error {
+	c := Container{
+		Container: step.Container,
+		Image:     step.Uses.RefID(),
+		Replicas:  replicas,
+	}
+
+	c.Envs = c.Envs.Merge(d.container.Envs)
+
+	c.Envs["PIPELINE_SCOPE"] = scope
+	c.Envs["PIPELINE_STAGE"] = stage
+
+	return d.podController.Apply(context.Background(), PodNameByScopeAndStage(scope, stage), &c)
+}
+
+func (d operatorMgr) Destroy(scope string, stage string) error {
+	return d.podController.Kill(context.Background(), PodNameByScopeAndStage(scope, stage))
+}

pipeline/operator/container/operatormgr_test.go

@@ -0,0 +1,64 @@
+package container
+
+import (
+	"testing"
+	"time"
+
+	"github.com/go-courier/semver"
+	"github.com/querycap/pipeline/spec"
+	"github.com/sirupsen/logrus"
+)
+
+func init() {
+	logrus.SetLevel(logrus.DebugLevel)
+}
+
+var scope, stage = "xxx", "fetch-image"
+var imageRegistry, _ = ParseImageRegistry("registry://docker.io/library/")
+
+func Test(t *testing.T) {
+	t.Run("docker", func(t *testing.T) {
+		c, _ := NewDockerPodController(imageRegistry)
+		RunAll(t, c)
+	})
+}
+
+func RunAll(t *testing.T, c PodController) {
+	s := spec.Stage{}
+	s.Uses.Name = "nginx"
+	s.Uses.Version = *semver.MustParseVersion("1.17.10")
+
+	mgr := NewOperatorMgr(c, &Container{})
+
+	t.Run("start", func(t *testing.T) {
+		if err := mgr.Up(scope, stage, s, 3); err != nil {
+			panic(err)
+		}
+
+		time.Sleep(1 * time.Second)
+
+		if err := mgr.Up(scope, stage, s, 5); err != nil {
+			panic(err)
+		}
+
+		time.Sleep(1 * time.Second)
+
+		if err := mgr.Up(scope, stage, s, 3); err != nil {
+			panic(err)
+		}
+
+		time.Sleep(1 * time.Second)
+
+		if err := mgr.Up(scope, stage, s, 0); err != nil {
+			panic(err)
+		}
+
+		time.Sleep(1 * time.Second)
+	})
+
+	t.Run("stop", func(t *testing.T) {
+		if err := mgr.Destroy(scope, stage); err != nil {
+			panic(err)
+		}
+	})
+}

pipeline/operator/container/pod_controller.go

@@ -0,0 +1,24 @@
+package container
+
+import (
+	"context"
+	"strings"
+
+	"github.com/querycap/pipeline/spec"
+)
+
+func PodNameByScopeAndStage(scope string, stage string) string {
+	return strings.Replace(scope+"--"+stage, ":", "__", -1)
+}
+
+type PodController interface {
+	Apply(ctx context.Context, name string, container *Container) error
+	Kill(ctx context.Context, name string) error
+}
+
+type Container struct {
+	spec.Container
+	Image       string
+	Replicas    int
+	Annotations map[string]string
+}

pipeline/operator/container/pod_controller__docker.go

@@ -0,0 +1,179 @@
+package container
+
+import (
+	"context"
+	"io"
+	"math/rand"
+	"os"
+	"strings"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/filters"
+	"github.com/docker/docker/client"
+	"github.com/docker/docker/pkg/stdcopy"
+	"github.com/sirupsen/logrus"
+)
+
+func NewDockerPodController(imageRegistry *ImageRegistry) (*DockerPodController, error) {
+	c, err := client.NewEnvClient()
+	if err != nil {
+		return nil, err
+	}
+
+	return &DockerPodController{
+		client:        c,
+		imageRegistry: imageRegistry,
+	}, nil
+}
+
+type DockerPodController struct {
+	client        *client.Client
+	imageRegistry *ImageRegistry
+}
+
+func (c *DockerPodController) Apply(ctx context.Context, name string, container *Container) error {
+	imageRef := c.imageRegistry.Fix(container.Image)
+
+	if err := c.PullImageIfNotExists(ctx, imageRef); err != nil {
+		return err
+	}
+
+	container.Annotations = map[string]string{
+		"pipeline": name,
+	}
+
+	list, err := c.ListMatchedRunningContainer(ctx, name)
+	if err != nil {
+		return err
+	}
+
+	current := len(list)
+
+	offset := current - container.Replicas
+
+	// scale up
+	if offset < 0 {
+		for i := 0; i < -offset; i++ {
+			if err := c.RunContainer(ctx, imageRef, container); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+
+	// scale down
+	if offset > 0 {
+		rand.Shuffle(len(list), func(i, j int) { list[i], list[j] = list[j], list[i] })
+
+		for _, item := range list[0:offset] {
+			if err := c.KillContainer(ctx, item.ID); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func (c *DockerPodController) Kill(ctx context.Context, name string) error {
+	list, err := c.ListMatchedRunningContainer(ctx, name)
+	if err != nil {
+		return err
+	}
+
+	for i := range list {
+		if err := c.KillContainer(ctx, list[i].ID); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (c *DockerPodController) ListMatchedRunningContainer(ctx context.Context, name string) ([]types.Container, error) {
+	containerListFilters := filters.NewArgs()
+	containerListFilters.Add("label", "pipeline="+name)
+
+	return c.ListRunningContainer(ctx, containerListFilters)
+}
+
+func (c *DockerPodController) RunContainer(ctx context.Context, image string, cc *Container) error {
+	logrus.WithContext(ctx).Debugf("running from %s", image)
+
+	containerConfig := &container.Config{
+		Image:  image,
+		Labels: cc.Annotations,
+	}
+
+	for k := range cc.Envs {
+		containerConfig.Env = append(containerConfig.Env, k+"="+cc.Envs[k])
+	}
+
+	created, err := c.client.ContainerCreate(ctx, containerConfig, nil, nil, "")
+	if err != nil {
+		return err
+	}
+
+	if err := c.client.ContainerStart(ctx, created.ID, types.ContainerStartOptions{}); err != nil {
+		return err
+	}
+
+	go func() {
+		if !logrus.IsLevelEnabled(logrus.DebugLevel) {
+			return
+		}
+
+		// clone logs
+		r, err := c.client.ContainerLogs(ctx, created.ID, types.ContainerLogsOptions{
+			ShowStdout: true,
+			ShowStderr: true,
+			Follow:     true,
+		})
+		if err != nil {
+			return
+		}
+		defer r.Close()
+
+		if _, err := stdcopy.StdCopy(os.Stdout, os.Stderr, r); err != nil && err != io.EOF {
+			logrus.Fatal(err)
+		}
+	}()
+
+	return nil
+}
+
+func (c *DockerPodController) ListRunningContainer(ctx context.Context, args filters.Args) ([]types.Container, error) {
+	return c.client.ContainerList(ctx, types.ContainerListOptions{
+		Filters: args,
+		All:     false,
+	})
+}
+
+func (c *DockerPodController) KillContainer(ctx context.Context, containerID string) error {
+	logrus.WithContext(ctx).Debugf("killing %s", containerID)
+
+	return c.client.ContainerKill(ctx, containerID, "SIGKILL")
+}
+
+func (c *DockerPodController) PullImageIfNotExists(ctx context.Context, image string) error {
+	imageListFilters := filters.NewArgs()
+	imageListFilters.Add("reference", strings.Replace(image, "docker.io/library/", "", 1))
+
+	images, err := c.client.ImageList(ctx, types.ImageListOptions{Filters: imageListFilters})
+	if err != nil {
+		return err
+	}
+
+	imageExists := len(images) > 0
+
+	if !imageExists {
+		if _, err = c.client.ImagePull(ctx, image, types.ImagePullOptions{
+			RegistryAuth: c.imageRegistry.RegistryAuth(),
+		}); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}