entry.go

// Copyright 2021 The Cockroach Authors.
//
// Licensed as a CockroachDB Enterprise file under the Cockroach Community
// License (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
//     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt

package tenant

import (
	"context"
	"fmt"
	"math/rand"
	"time"

	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/cockroachdb/cockroach/pkg/util/randutil"
	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
	"github.com/cockroachdb/errors"
)

// tenantEntry is an entry in the tenant directory that records information
// about a single tenant, including its ID, cluster name, and the IP addresses
// for available pods.
type tenantEntry struct {
	// These fields can be read by callers without synchronization, since
	// they're written once during initialization, and are immutable thereafter.

	// TenantID is the identifier for this tenant which is unique within a CRDB
	// cluster.
	TenantID roachpb.TenantID

	// Full name of the tenant's cluster i.e. dim-dog.
	ClusterName string

	// RefreshDelay is the minimum amount of time that must elapse between
	// attempts to refresh pods for this tenant after ReportFailure is
	// called.
	RefreshDelay time.Duration

	// initialized is set to true once Initialized has been called.
	initialized bool

	// initError is set to any error that occurs in Initialized (or nil if no
	// error occurred).
	initError error

	// pods synchronizes access to information about the tenant's SQL
	// pods. These fields can be updated over time, so a lock must be
	// obtained before accessing them.
	pods struct {
		syncutil.Mutex
		rng  *rand.Rand
		pods []tenantPod
	}

	// calls synchronizes calls to the K8s API for this tenant (e.g. calls to
	// RefreshPods). Synchronization is needed to ensure that only one
	// thread at a time is calling on behalf of a tenant, and that calls are rate
	// limited to prevent storms.
	calls struct {
		syncutil.Mutex
		lastRefresh time.Time
	}
}

// Initialize fetches metadata about a tenant, such as its cluster name, and
// stores that in the entry. After this is called once, all future calls return
// the same result (and do nothing).
func (e *tenantEntry) Initialize(ctx context.Context, client DirectoryClient) error {
	// Synchronize multiple threads trying to initialize. Only the first thread
	// does the initialization.
	e.calls.Lock()
	defer e.calls.Unlock()

	// If Initialize has already been called, return any error that occurred.
	if e.initialized {
		return e.initError
	}

	tenantResp, err := client.GetTenant(ctx, &GetTenantRequest{TenantID: e.TenantID.ToUint64()})
	if err != nil {
		e.initialized = true
		e.initError = err
		return err
	}

	e.ClusterName = tenantResp.ClusterName

	e.pods.rng, _ = randutil.NewPseudoRand()

	e.initialized = true
	return nil
}

// RefreshPods makes a synchronous directory server call to fetch the
// latest information about the tenant's available pods, such as their IP
// addresses.
func (e *tenantEntry) RefreshPods(ctx context.Context, client DirectoryClient) error {
	if !e.initialized {
		return errors.AssertionFailedf("entry for tenant %d is not initialized", e.TenantID)
	}

	// Lock so that only one thread at a time will refresh, since there's no
	// point in multiple threads doing it within a short span of time - it's
	// likely nothing has changed.
	e.calls.Lock()
	defer e.calls.Unlock()

	// If refreshed recently, no-op.
	if !e.canRefreshLocked() {
		return nil
	}

	log.Infof(ctx, "refreshing tenant %d pods", e.TenantID)

	_, err := e.fetchPodsLocked(ctx, client)
	return err
}

// ChoosePodAddr returns the IP address of one of this tenant's available
// pods. If a tenant has multiple pods, then ChoosePodAddr returns
// the IP address of one of those pods. If the tenant is suspended and no
// pods are available, then ChoosePodAddr will trigger resumption of the
// tenant and return the IP address of the new pod. Note that resuming a
// tenant requires directory server calls, so ChoosePodAddr can block for
// some time, until the resumption process is complete. However, if
// errorIfNoPods is true, then ChoosePodAddr returns an error if there
// are no pods available rather than blocking.
func (e *tenantEntry) ChoosePodAddr(
	ctx context.Context, client DirectoryClient, errorIfNoPods bool,
) (string, error) {
	if !e.initialized {
		return "", errors.AssertionFailedf("entry for tenant %d is not initialized", e.TenantID)
	}

	pods := e.getPods()
	if len(pods) == 0 {
		// There are no known pod IP addresses, so fetch pod information
		// from the directory server. Resume the tenant if it is suspended; that
		// will always result in at least one pod IP address (or an error).
		var err error
		if pods, err = e.ensureTenantPod(ctx, client, errorIfNoPods); err != nil {
			return "", err
		}
	}

	e.pods.Lock()
	defer e.pods.Unlock()

	return selectTenantPod(e.pods.rng, pods).Addr, nil
}

// AddPodAddr inserts the given IP address and load into the tenant's list of
// pods. If it is already present, then .Load is updated and AddPodAddr returns
// false.
func (e *tenantEntry) AddPodAddr(addr string, load float32) bool {
	e.pods.Lock()
	defer e.pods.Unlock()

	for i, existing := range e.pods.pods {
		if existing.Addr == addr {
			e.pods.pods[i].Load = load
			return false
		}
	}

	e.pods.pods = append(e.pods.pods, tenantPod{Addr: addr, Load: load})
	return true
}

// RemovePodAddr removes the given IP address from the tenant's list of
// pod addresses. If it was not present, RemovePodAddr returns false.
func (e *tenantEntry) RemovePodAddr(addr string) bool {
	e.pods.Lock()
	defer e.pods.Unlock()

	for i, existing := range e.pods.pods {
		if existing.Addr == addr {
			copy(e.pods.pods[i:], e.pods.pods[i+1:])
			e.pods.pods = e.pods.pods[:len(e.pods.pods)-1]
			return true
		}
	}
	return false
}

// getPod gets the current list of pods within scope of lock and returns them.
func (e *tenantEntry) getPods() []tenantPod {
	e.pods.Lock()
	defer e.pods.Unlock()
	return e.pods.pods
}

// ensureTenantPod ensures that at least one SQL process exists for this
// tenant, and is ready for connection attempts to its IP address. If
// errorIfNoPods is true, then ensureTenantPod returns an error if
// there are no pods available rather than blocking.
func (e *tenantEntry) ensureTenantPod(
	ctx context.Context, client DirectoryClient, errorIfNoPods bool,
) (pods []tenantPod, err error) {
	const retryDelay = 100 * time.Millisecond

	e.calls.Lock()
	defer e.calls.Unlock()

	// If an IP address is already available, nothing more to do. Check this
	// immediately after obtaining the lock so that only the first thread does
	// the work to get information about the tenant.
	pods = e.getPods()
	if len(pods) != 0 {
		return pods, nil
	}

	for {
		// Check for context cancellation or timeout.
		if err = ctx.Err(); err != nil {
			return nil, err
		}

		// Try to resume the tenant if not yet resumed.
		_, err = client.EnsurePod(ctx, &EnsurePodRequest{e.TenantID.ToUint64()})
		if err != nil {
			return nil, err
		}

		// Get pod information for the newly resumed tenant. Except in rare
		// race conditions, this is expected to immediately find an IP address,
		// since the above call started a tenant process that already has an IP
		// address.
		pods, err = e.fetchPodsLocked(ctx, client)
		if err != nil {
			return nil, err
		}
		if len(pods) != 0 {
			log.Infof(ctx, "resumed tenant %d", e.TenantID)
			break
		}

		// In rare case where no IP address is ready, wait for a bit before
		// retrying.
		if errorIfNoPods {
			return nil, fmt.Errorf("no pods available for tenant %s", e.TenantID)
		}
		sleepContext(ctx, retryDelay)
	}

	return pods, nil
}

// fetchPodsLocked makes a synchronous directory server call to get the
// latest information about the tenant's available pods, such as their IP
// addresses.
//
// NOTE: Caller must lock the "calls" mutex before calling fetchPodsLocked.
func (e *tenantEntry) fetchPodsLocked(
	ctx context.Context, client DirectoryClient,
) (tenantPods []tenantPod, err error) {
	// List the pods for the given tenant.
	list, err := client.ListPods(ctx, &ListPodsRequest{e.TenantID.ToUint64()})
	if err != nil {
		return nil, err
	}

	// Get updated list of running process pod IP addresses and save it to
	// the entry.
	tenantPods = make([]tenantPod, 0, len(list.Pods))
	for i := range list.Pods {
		pod := list.Pods[i]
		tenantPods = append(tenantPods, tenantPod{
			Addr: pod.Addr,
			// Load: pod.Load,
		})
	}

	// Need to lock in case another thread is reading the IP addresses (e.g. in
	// ChoosePodAddr).
	e.pods.Lock()
	defer e.pods.Unlock()
	e.pods.pods = tenantPods

	if len(tenantPods) != 0 {
		log.Infof(ctx, "fetched IP addresses: %v", tenantPods)
	}

	return tenantPods, nil
}

// canRefreshLocked returns true if it's been at least X milliseconds since the
// last time the tenant pod information was refreshed. This has the effect
// of rate limiting RefreshPods calls.
//
// NOTE: Caller must lock the "calls" mutex before calling canRefreshLocked.
func (e *tenantEntry) canRefreshLocked() bool {
	now := timeutil.Now()
	if now.Sub(e.calls.lastRefresh) < e.RefreshDelay {
		return false
	}
	e.calls.lastRefresh = now
	return true
}