egress-proxies-operator/internal/provider/kvm2/hosts.go

package kvm2

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/digitalocean/go-libvirt"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	logf "sigs.k8s.io/controller-runtime/pkg/log"

	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/provider"
	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/tracing"
)

// dialer opens one connection to a host. The production dialer is
// libvirtDialer; tests inject one that hands out fakes.
type dialer interface {
	dial(ctx context.Context, host provider.KVMHost) (backend, error)
}

// hostTable holds the pooled connection per host (spec entry 19). One dial
// per host at a time: concurrent callers wait for it and share the result.
// A dead backend (abandoned call, keepalive failure, socket lost) is dropped
// and redialled on the next get.
//
// What a cancelled or overlong call does to the shared connection (spec
// entry 18): the adapter kills it, so every other call in flight on that
// host fails once with Transient and its caller retries on a fresh
// connection. Callers that must not be exposed to that, the image transfer
// and the Checker, use dialFresh and get a connection of their own.
type hostTable struct {
	dialer dialer

	mu    sync.Mutex
	hosts map[string]*hostEntry
}

// hostEntry serialises dials for one host. The semaphore, not a mutex, so a
// waiter can give up when its ctx ends.
type hostEntry struct {
	sem     chan struct{}
	backend backend
}

func newHostTable(d dialer) *hostTable {
	return &hostTable{dialer: d, hosts: make(map[string]*hostEntry)}
}

func (t *hostTable) entry(name string) *hostEntry {
	t.mu.Lock()
	defer t.mu.Unlock()
	e, ok := t.hosts[name]
	if !ok {
		e = &hostEntry{sem: make(chan struct{}, 1)}
		t.hosts[name] = e
	}
	return e
}

// get returns the live pooled backend for host, dialling it when there is
// none. Waiting for another caller's dial ends with ctx.
func (t *hostTable) get(ctx context.Context, host provider.KVMHost) (backend, error) {
	e := t.entry(host.Name)
	select {
	case e.sem <- struct{}{}:
	case <-ctx.Done():
		return nil, fmt.Errorf("%w: waiting for the connection to host %s: %w", provider.ErrTransient, host.Name, ctx.Err())
	}
	defer func() { <-e.sem }()
	if e.backend != nil {
		if !isDead(e.backend) {
			return e.backend, nil
		}
		_ = e.backend.Close()
		e.backend = nil
	}
	b, err := t.dialer.dial(ctx, host)
	if err != nil {
		return nil, err
	}
	e.backend = b
	return b, nil
}

// dialFresh returns a dedicated connection that is never pooled. The
// caller owns it and must Close it.
func (t *hostTable) dialFresh(ctx context.Context, host provider.KVMHost) (backend, error) {
	return t.dialer.dial(ctx, host)
}

// closeAll closes every pooled connection, waiting for in-flight dials.
func (t *hostTable) closeAll() error {
	t.mu.Lock()
	entries := make([]*hostEntry, 0, len(t.hosts))
	for _, e := range t.hosts {
		entries = append(entries, e)
	}
	t.mu.Unlock()
	var errs []error
	for _, e := range entries {
		e.sem <- struct{}{}
		if e.backend != nil {
			errs = append(errs, e.backend.Close())
			e.backend = nil
		}
		<-e.sem
	}
	return errors.Join(errs...)
}

// libvirtDialer is the production dialer: SSH transport, go-libvirt over
// the forwarded socket, the connect handshake, then the observe decorator,
// so provider code never holds an undecorated backend.
type libvirtDialer struct {
	providerName string
	cfg          provider.KVMConfig
	budgets      budgets
	keepalive    time.Duration
	tracing      []tracing.Option
}

func newLibvirtDialer(providerName string, cfg provider.KVMConfig) *libvirtDialer {
	return &libvirtDialer{
		providerName: providerName,
		cfg:          cfg,
		budgets:      defaultBudgets,
		keepalive:    defaultKeepaliveInterval,
	}
}

func (d *libvirtDialer) dial(ctx context.Context, host provider.KVMHost) (backend, error) {
	ctx, span := tracing.StartSpan(ctx, tracing.Tracer(d.tracing...), "kvm2.dial",
		trace.WithSpanKind(trace.SpanKindClient),
		trace.WithAttributes(
			attribute.String("kvm2.host", host.Name),
			attribute.String("kvm2.address", host.Address),
			attribute.String("kvm2.user", host.User),
		))
	defer span.End()
	log := logf.FromContext(ctx).WithName("kvm2").WithValues("provider", d.providerName, "host", host.Name)
	start := time.Now()
	b, err := d.connect(ctx, host)
	if err != nil {
		span.SetStatus(codes.Error, err.Error())
		kv := []any{"address", host.Address, "user", host.User, "duration", time.Since(start), "error", err.Error()}
		if stage, ok := stageOf(err); ok {
			kv = append(kv, "stage", string(stage))
		}
		kv = append(kv, libvirtErrorFields(err)...)
		log.V(1).Info("libvirt host connect failed", kv...)
		return nil, err
	}
	log.V(1).Info("libvirt host connected", "address", host.Address, "user", host.User, "duration", time.Since(start))
	return observe(b, d.providerName, host.Name, d.tracing...), nil
}

// connect is the dial sequence under one budget: the SSH stages, then the
// libvirt handshake. On success the keepalive and the disconnect watcher
// run until the link dies.
func (d *libvirtDialer) connect(ctx context.Context, host provider.KVMHost) (*libvirtBackend, error) {
	ctx, cancel := context.WithTimeout(ctx, d.budgets.dial)
	defer cancel()
	sock, err := dialSSHSocket(ctx, host, d.cfg, 0)
	if err != nil {
		return nil, err
	}
	b := newLibvirtBackend(libvirt.New(sock), host.Name, sock, d.budgets)
	if err := b.connect(ctx); err != nil {
		b.link.markDead("connect failed")
		return nil, fmt.Errorf("connecting to libvirt on host %s: %w", host.Name, err)
	}
	go func() {
		if err := sock.keepalive(d.keepalive, b.link.dead()); err != nil {
			b.link.markDead(err.Error())
		}
	}()
	return b, nil
}