docker-compose/local/dependencies.go

// +build local

/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package local

import (
	"context"
	"fmt"
	"strings"
	"sync"

	"github.com/compose-spec/compose-go/types"
	"golang.org/x/sync/errgroup"
)

// ServiceStatus indicates the status of a service
type ServiceStatus int

// Services status flags
const (
	ServiceStopped ServiceStatus = iota
	ServiceStarted
)

func inDependencyOrder(ctx context.Context, project *types.Project, fn func(context.Context, types.ServiceConfig) error) error {
	g := NewGraph(project.Services)
	if b, err := g.HasCycles(); b {
		return err
	}

	leaves := g.Leaves()

	eg, _ := errgroup.WithContext(ctx)
	eg.Go(func() error {
		return run(ctx, g, eg, leaves, fn)
	})

	return eg.Wait()
}

// Note: this could be `graph.walk` or whatever
func run(ctx context.Context, graph *Graph, eg *errgroup.Group, nodes []*Vertex, fn func(context.Context, types.ServiceConfig) error) error {
	for _, node := range nodes {
		n := node
		// Don't start this service yet if all of its children have
		// not been started yet.
		if len(graph.FilterChildren(n.Service.Name, ServiceStopped)) != 0 {
			continue
		}

		eg.Go(func() error {
			err := fn(ctx, n.Service)
			if err != nil {
				return err
			}

			graph.UpdateStatus(n.Service.Name, ServiceStarted)

			return run(ctx, graph, eg, n.GetParents(), fn)
		})
	}

	return nil
}

// Graph represents project as service dependencies
type Graph struct {
	Vertices map[string]*Vertex
	lock     sync.RWMutex
}

// Vertex represents a service in the dependencies structure
type Vertex struct {
	Key      string
	Service  types.ServiceConfig
	Status   ServiceStatus
	Children map[string]*Vertex
	Parents  map[string]*Vertex
}

// GetParents returns a slice with the parent vertexes of the current Vertex
func (v *Vertex) GetParents() []*Vertex {
	var res []*Vertex
	for _, p := range v.Parents {
		res = append(res, p)
	}
	return res
}

// NewGraph returns the dependency graph of the services
func NewGraph(services types.Services) *Graph {
	graph := &Graph{
		lock:     sync.RWMutex{},
		Vertices: map[string]*Vertex{},
	}

	for _, s := range services {
		graph.AddVertex(s.Name, s)
	}

	for _, s := range services {
		for _, name := range s.GetDependencies() {
			_ = graph.AddEdge(s.Name, name)
		}
	}

	return graph
}

// NewVertex is the constructor function for the Vertex
func NewVertex(key string, service types.ServiceConfig) *Vertex {
	return &Vertex{
		Key:      key,
		Service:  service,
		Status:   ServiceStopped,
		Parents:  map[string]*Vertex{},
		Children: map[string]*Vertex{},
	}
}

// AddVertex adds a vertex to the Graph
func (g *Graph) AddVertex(key string, service types.ServiceConfig) {
	g.lock.Lock()
	defer g.lock.Unlock()

	v := NewVertex(key, service)
	g.Vertices[key] = v
}

// AddEdge adds a relationship of dependency between vertexes `source` and `destination`
func (g *Graph) AddEdge(source string, destination string) error {
	g.lock.Lock()
	defer g.lock.Unlock()

	sourceVertex := g.Vertices[source]
	destinationVertex := g.Vertices[destination]

	if sourceVertex == nil {
		return fmt.Errorf("could not find %s", source)
	}
	if destinationVertex == nil {
		return fmt.Errorf("could not find %s", destination)
	}

	// If they are already connected
	if _, ok := sourceVertex.Children[destination]; ok {
		return nil
	}

	sourceVertex.Children[destination] = destinationVertex
	destinationVertex.Parents[source] = sourceVertex

	return nil
}

// Leaves returns the slice of leaves of the graph
func (g *Graph) Leaves() []*Vertex {
	g.lock.Lock()
	defer g.lock.Unlock()

	var res []*Vertex
	for _, v := range g.Vertices {
		if len(v.Children) == 0 {
			res = append(res, v)
		}
	}

	return res
}

// UpdateStatus updates the status of a certain vertex
func (g *Graph) UpdateStatus(key string, status ServiceStatus) {
	g.lock.Lock()
	defer g.lock.Unlock()
	g.Vertices[key].Status = status
}

// FilterChildren returns children of a certain vertex that are in a certain status
func (g *Graph) FilterChildren(key string, status ServiceStatus) []*Vertex {
	g.lock.Lock()
	defer g.lock.Unlock()

	var res []*Vertex
	vertex := g.Vertices[key]

	for _, child := range vertex.Children {
		if child.Status == status {
			res = append(res, child)
		}
	}

	return res
}

// HasCycles detects cycles in the graph
func (g *Graph) HasCycles() (bool, error) {
	discovered := []string{}
	finished := []string{}

	for _, vertex := range g.Vertices {
		path := []string{
			vertex.Key,
		}
		if !contains(discovered, vertex.Key) && !contains(finished, vertex.Key) {
			var err error
			discovered, finished, err = g.visit(vertex.Key, path, discovered, finished)

			if err != nil {
				return true, err
			}
		}
	}

	return false, nil
}

func (g *Graph) visit(key string, path []string, discovered []string, finished []string) ([]string, []string, error) {
	discovered = append(discovered, key)

	for _, v := range g.Vertices[key].Children {
		path := append(path, v.Key)
		if contains(discovered, v.Key) {
			return nil, nil, fmt.Errorf("cycle found: %s", strings.Join(path, " -> "))
		}

		if !contains(finished, v.Key) {
			if _, _, err := g.visit(v.Key, path, discovered, finished); err != nil {
				return nil, nil, err
			}
		}
	}

	discovered = remove(discovered, key)
	finished = append(finished, key)
	return discovered, finished, nil
}

func remove(slice []string, item string) []string {
	var s []string
	for _, i := range slice {
		if i != item {
			s = append(s, i)
		}
	}
	return s
}