entry_manager.go

/*
Copyright 2019 The Skaffold 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 portforward

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/color"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event"
	"k8s.io/apimachinery/pkg/util/wait"
)

var (
	// For testing
	forwardingTimeoutTime = time.Minute
	portForwardEvent      = func(entry *portForwardEntry) {
		// TODO priyawadhwa@, change event API to accept ports of type int
		event.PortForwarded(
			int32(entry.localPort),
			int32(entry.resource.Port),
			entry.podName,
			entry.containerName,
			entry.resource.Namespace,
			entry.portName,
			string(entry.resource.Type),
			entry.resource.Name)
	}
)

type forwardedPorts struct {
	ports map[int]struct{}
	lock  *sync.Mutex
}

func (f forwardedPorts) Store(key, _ interface{}) {
	f.lock.Lock()
	defer f.lock.Unlock()
	val, ok := key.(int)
	if !ok {
		panic("only store keys of type int in forwardedPorts")
	}
	// this map is only used as a set of keys, we don't care about the values
	f.ports[val] = dummy()
}

func (f forwardedPorts) LoadOrStore(key, _ interface{}) (interface{}, bool) {
	f.lock.Lock()
	defer f.lock.Unlock()
	k, ok := key.(int)
	if !ok {
		return nil, false
	}
	// this map is only used as a set of keys, we don't care about the values
	_, exists := f.ports[k]
	return dummy(), exists
}

func dummy() struct{} {
	return struct{}{}
}

func (f forwardedPorts) Delete(port int) {
	f.lock.Lock()
	defer f.lock.Unlock()
	delete(f.ports, port)
}

type forwardedResources struct {
	resources map[string]*portForwardEntry
	lock      *sync.Mutex
}

func (f forwardedResources) Store(key, value interface{}) {
	f.lock.Lock()
	defer f.lock.Unlock()
	k, ok := key.(string)
	if !ok {
		panic("only store keys of type string in forwardedResources")
	}
	val, ok := value.(*portForwardEntry)
	if !ok {
		panic("only store values of type *portForwardEntry in forwardedResources")
	}
	f.resources[k] = val
}

func (f forwardedResources) Load(key string) (*portForwardEntry, bool) {
	f.lock.Lock()
	defer f.lock.Unlock()
	val, exists := f.resources[key]
	return val, exists
}

func (f forwardedResources) Delete(resource string) {
	f.lock.Lock()
	defer f.lock.Unlock()
	delete(f.resources, resource)
}

func (f forwardedResources) Length() int {
	f.lock.Lock()
	defer f.lock.Unlock()
	return len(f.resources)
}

func newForwardedPorts() forwardedPorts {
	return forwardedPorts{
		lock:  &sync.Mutex{},
		ports: map[int]struct{}{},
	}
}

func newForwardedResources() forwardedResources {
	return forwardedResources{
		lock:      &sync.Mutex{},
		resources: map[string]*portForwardEntry{},
	}
}

// EntryManager handles forwarding entries and keeping track of
// forwarded ports and resources.
type EntryManager struct {
	EntryForwarder
	output io.Writer

	// forwardedPorts serves as a synchronized set of ports we've forwarded.
	forwardedPorts forwardedPorts

	// forwardedResources is a map of portForwardEntry key (string) -> portForwardEntry
	forwardedResources forwardedResources
}

// NewEntryManager returns a new port forward entry manager to keep track
// of forwarded ports and resources
func NewEntryManager(out io.Writer) EntryManager {
	return EntryManager{
		output:             out,
		forwardedPorts:     newForwardedPorts(),
		forwardedResources: newForwardedResources(),
		EntryForwarder:     &KubectlForwarder{},
	}
}

func (b *EntryManager) forwardPortForwardEntry(ctx context.Context, entry *portForwardEntry) error {
	// Check if this resource has already been forwarded
	if _, ok := b.forwardedResources.Load(entry.key()); ok {
		return nil
	}
	b.forwardedResources.Store(entry.key(), entry)
	err := wait.PollImmediate(time.Second, forwardingTimeoutTime, func() (bool, error) {
		if err := b.Forward(ctx, entry); err != nil {
			return false, nil
		}
		return true, nil
	})

	go b.Monitor(entry, func() {
		b.Retry(ctx, entry)
	})

	if err != nil {
		return err
	}

	color.Default.Fprintln(b.output, fmt.Sprintf("Port forwarded %s/%s from remote port %d to local port %d", entry.resource.Type, entry.resource.Name, entry.resource.Port, entry.localPort))

	portForwardEvent(entry)
	return nil
}

// Stop terminates all kubectl port-forward commands.
func (b *EntryManager) Stop() {
	for _, pfe := range b.forwardedResources.resources {
		b.Terminate(pfe)
	}
}

// Terminate terminates a single port forward entry
func (b *EntryManager) Terminate(p *portForwardEntry) {
	b.forwardedResources.Delete(p.key())
	b.forwardedPorts.Delete(p.localPort)
	b.EntryForwarder.Terminate(p)
}

func (b *EntryManager) Retry(ctx context.Context, p *portForwardEntry) error {
	b.Terminate(p)
	return b.forwardPortForwardEntry(ctx, p)
}