prom-autoexport.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"os"
	"text/template"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/events"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/docker/client"
	"github.com/sirupsen/logrus"
)

type Exporter struct {
	Image    string
	Cmd      string
	Exported types.ContainerJSON
}

const (
	LABEL_EXPORTED_ID    = "autoexporter.exported.id"
	LABEL_EXPORTED_NAME  = "autoexporter.exported.name"
	LABEL_ENABLE         = "autoexporter.enable"
	LABEL_EXPORTER_IMAGE = "autoexporter.exporter"
	LABEL_EXPORTER_CMD   = "autoexporter.cmd"
)

func run() error {
	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithVersion("1.37"))
	if err != nil {
		return err
	}
	defer cli.Close()

	ctx := context.Background()

	if err := cleanup(ctx, cli); err != nil {
		return err
	}

	evtCh, errCh := cli.Events(ctx, types.EventsOptions{
		Since: time.Now().Format(time.RFC3339),
	})

	logrus.Info("Start listening for new Docker events...")

	for {
		select {
		case err := <-errCh:
			// @TODO: should auto-reconnect instead
			panic(err)
		case evt := <-evtCh:
			go func() {
				if err := handleEvent(ctx, cli, evt); err != nil {
					logrus.Error(err)
				}
			}()
		}
	}

	return nil
}

func cleanup(ctx context.Context, cli *client.Client) error {
	containers, err := cli.ContainerList(ctx, types.ContainerListOptions{
		Filters: filters.NewArgs(filters.KeyValuePair{
			Key:   "label",
			Value: LABEL_EXPORTED_ID,
		}),
	})

	if err != nil {
		return err
	}

	for _, container := range containers {
		go func(container types.Container) {
			if err := cleanupExporter(ctx, cli, container); err != nil {
				logrus.Error(err)
			}
		}(container)
	}

	return nil
}

func cleanupExporter(ctx context.Context, cli *client.Client, exporter types.Container) error {
	_, err := cli.ContainerInspect(ctx, exporter.Labels[LABEL_EXPORTED_NAME])

	// If the exported service is still alive, we stop cleanup process
	if err == nil {
		return nil
	}
	if !client.IsErrNotFound(err) {
		return err
	}

	stopExporter(ctx, cli, exporter.Labels[LABEL_EXPORTED_NAME])

	return nil
}

func handleEvent(ctx context.Context, cli *client.Client, evt events.Message) error {
	if evt.Type == "container" && evt.Action == "start" {
		return handleContainerStart(ctx, cli, evt)
	} else if evt.Type == "container" && evt.Action == "stop" {
		return handleContainerStop(ctx, cli, evt)
	}

	return nil
}

func handleContainerStart(ctx context.Context, cli *client.Client, evt events.Message) error {
	container, err := cli.ContainerInspect(ctx, evt.Actor.ID)
	if err != nil {
		return err
	}

	enable, ok := container.Config.Labels[LABEL_ENABLE]
	if !ok || !strToBool(enable) {
		return nil
	}

	exporterImg, err := readLabel(container, LABEL_EXPORTER_IMAGE)
	if err != nil {
		return err
	}

	exporterCmd, err := readLabel(container, LABEL_EXPORTER_CMD)
	if err != nil {
		return err
	}

	exporter := Exporter{
		Image:    exporterImg,
		Cmd:      exporterCmd,
		Exported: container,
	}

	logrus.WithFields(logrus.Fields{
		LABEL_EXPORTER_IMAGE: exporter.Image,
		LABEL_EXPORTER_CMD:   exporter.Cmd,
	}).Infof("Container %q started.", exporter.Exported.Name, exporter.Image, exporter.Cmd)

	return runExporter(ctx, cli, exporter)
}

func strToBool(str string) bool {
	return str == "yes" || str == "true" || str == "on" || str == "1"
}

func readLabel(container types.ContainerJSON, label string) (string, error) {
	labelVal := container.Config.Labels[label]
	tpl, err := template.New(label).Parse(labelVal)
	if err != nil {
		return "", err
	}

	var buf bytes.Buffer
	writer := bufio.NewWriter(&buf)
	err = tpl.Execute(writer, container)
	if err != nil {
		return "", err
	}

	writer.Flush()
	val := buf.String()

	return val, nil
}

func runExporter(ctx context.Context, cli *client.Client, exporter Exporter) error {
	config := container.Config{
		User:  "1000",
		Cmd:   strslice.StrSlice{exporter.Cmd},
		Image: exporter.Image,
		Labels: map[string]string{
			LABEL_EXPORTED_ID:   exporter.Exported.ID,
			LABEL_EXPORTED_NAME: exporter.Exported.Name,
		},
	}
	hostConfig := container.HostConfig{
		NetworkMode: container.NetworkMode(fmt.Sprintf("container:%s", exporter.Exported.ID)),
		RestartPolicy: container.RestartPolicy{
			Name:              "on-failure",
			MaximumRetryCount: 10,
		},
	}
	networkingConfig := network.NetworkingConfig{}

	container, err := cli.ContainerCreate(ctx, &config, &hostConfig, &networkingConfig, fmt.Sprintf("%s-exporter", exporter.Exported.Name))
	if err != nil {
		return err
	}

	if len(container.Warnings) > 0 {
		logrus.Warningf("Warnings during exporter creation: %+v", container.Warnings)
	}

	err = cli.ContainerStart(ctx, container.ID, types.ContainerStartOptions{})
	if err != nil {
		return err
	}

	logrus.WithFields(logrus.Fields{
		"exporter_cid": container.ID,
		"image":        exporter.Image,
		"cmd":          exporter.Cmd,
	}).Infof("Autoexporter for container %q started.", exporter.Exported.Name)

	return nil
}

func handleContainerStop(ctx context.Context, cli *client.Client, evt events.Message) error {
	container, err := cli.ContainerInspect(ctx, evt.Actor.ID)
	if err != nil {
		return err
	}

	enable, ok := container.Config.Labels["autoexporter.enable"]
	if !ok || !strToBool(enable) {
		return nil
	}

	return stopExporter(ctx, cli, container.Name)
}

func stopExporter(ctx context.Context, cli *client.Client, exported string) error {
	exporter := fmt.Sprintf("%s-exporter", exported)
	err := cli.ContainerStop(ctx, exporter, nil)
	if err != nil {
		return err
	}

	err = cli.ContainerRemove(ctx, exporter, types.ContainerRemoveOptions{
		Force: true,
	})
	if err != nil {
		return err
	}

	logrus.WithFields(logrus.Fields{
		"exporter": exporter,
		"exported": exported,
	}).Infof("Autoexporter for container %q stopped.", exported)

	return nil
}

func main() {
	if err := run(); err != nil {
		logrus.Fatal(err)
		os.Exit(1)
	}
}