<commit_before>package connector

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/pflag"
	v1 "k8s.io/api/core/v1"

	"github.com/datawire/ambassador/pkg/kates"
	"github.com/datawire/dlib/dexec"
	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dutil"
	"github.com/datawire/telepresence2/rpc/v2/daemon"
	"github.com/datawire/telepresence2/v2/pkg/client"
)

const connectTimeout = 5 * time.Second

// k8sCluster is a Kubernetes cluster reference
type k8sCluster struct {
	// Things for bypassing kates
	Server       string
	Context      string
	Namespace    string
	kubeFlagArgs []string

	// Main
	client *kates.Client
	daemon daemon.DaemonClient // for DNS updates

	// Current snapshot.
	// These fields (except for accLock/accCond) get set by acc.Update().
	accLock     sync.Mutex
	accCond     *sync.Cond
	Deployments []*kates.Deployment
	Pods        []*kates.Pod
	Services    []*kates.Service
}

// getKubectlArgs returns the kubectl command arguments to run a
// kubectl command with this cluster.
func (kc *k8sCluster) getKubectlArgs(args ...string) []string {
	return append(kc.kubeFlagArgs, args...)
}

// getKubectlCmd returns a Cmd that runs kubectl with the given arguments and
// the appropriate environment to talk to the cluster
func (kc *k8sCluster) getKubectlCmd(c context.Context, args ...string) *dexec.Cmd {
	return dexec.CommandContext(c, "kubectl", kc.getKubectlArgs(args...)...)
}

// portForwardAndThen starts a kubectl port-forward command, passes its output to the given scanner, and waits for
// the scanner to produce a result. A new goroutine is started as a sibling in the current dgroup.Group if the scanner
// returns something other than nil. The sibling is named after the argument thenName and executes the then function,
// passing what the scanner returned as an argument.
func (kc *k8sCluster) portForwardAndThen(
	c context.Context,
	kpfArgs []string,
	outputHandler func(*bufio.Scanner) interface{},
	thenName string,
	then func(context.Context, interface{}) error,
) error {
	args := make([]string, 0, len(kc.kubeFlagArgs)+1+len(kpfArgs))
	args = append(args, kc.kubeFlagArgs...)
	args = append(args, "port-forward")
	args = append(args, kpfArgs...)

	pf := dexec.CommandContext(c, "kubectl", args...)
	out, err := pf.StdoutPipe()
	if err != nil {
		return err
	}
	defer out.Close()

	// We want this command to keep on running. If it returns an error, then it was unsuccessful.
	if err = pf.Start(); err != nil {
		out.Close()
		dlog.Errorf(c, "port-forward failed to start: %v", client.RunError(err))
		return err
	}

	sc := bufio.NewScanner(out)

	// Give port-forward 10 seconds to produce the correct output and spawn the next process
	timer := time.AfterFunc(10*time.Second, func() {
		_ = pf.Process.Kill()
	})

	// wait group is done when next process starts.
	go func() {
		defer timer.Stop()
		if output := outputHandler(sc); output != nil {
			dgroup.ParentGroup(c).Go(thenName, func(c context.Context) error { return then(c, output) })
		}
	}()

	// let the port forward continue running. It will either be killed by the
	// timer (if it didn't produce the expected output) or by a context cancel.
	if err = pf.Wait(); err != nil {
		if c.Err() != nil {
			// Context cancelled
			err = nil
		} else {
			err = client.RunError(err)
			dlog.Errorf(c, "port-forward failed: %v", err)
		}
	}
	return err
}

// check for cluster connectivity
func (kc *k8sCluster) check(c context.Context) error {
	c, cancel := context.WithTimeout(c, connectTimeout)
	defer cancel()
	cmd := kc.getKubectlCmd(c, "get", "po", "ohai", "--ignore-not-found")
	stderr := bytes.Buffer{}
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		if c.Err() == context.DeadlineExceeded {
			err = errors.New("timeout when testing cluster connectivity")
		} else if msg := stderr.String(); len(msg) > 0 {
			err = errors.New(msg)
		}
	}
	return err
}

func (kc *k8sCluster) createWatch(c context.Context, namespace string) (acc *kates.Accumulator, err error) {
	defer func() {
		if r := dutil.PanicToError(recover()); r != nil {
			err = r
		}
	}()

	return kc.client.Watch(c,
		kates.Query{
			Name:      "Services",
			Namespace: namespace,
			Kind:      "service",
		},
		kates.Query{
			Name:      "Deployments",
			Namespace: namespace,
			Kind:      "deployment",
		},
		kates.Query{
			Name:      "Pods",
			Namespace: namespace,
			Kind:      "pod",
		}), nil
}

func (kc *k8sCluster) waitUntil(ctx context.Context, fn func(context.Context) (bool, error)) error {
	kc.accLock.Lock()
	defer kc.accLock.Unlock()
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		done, err := fn(ctx)
		if err != nil {
			return err
		}
		if done {
			return nil
		}
		kc.accCond.Wait()
	}
}

func (kc *k8sCluster) startWatches(c context.Context, namespace string, accWait chan<- struct{}) error {
	acc, err := kc.createWatch(c, namespace)
	if err != nil {
		return err
	}

	closeAccWait := func() {
		if accWait != nil {
			close(accWait)
			accWait = nil
		}
	}

	dgroup.ParentGroup(c).Go("watch-k8s", func(c context.Context) error {
		for {<commit_msg>Ensure that wait for accumulator condition ends on Context cancellation<commit_after>package connector

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/pflag"
	v1 "k8s.io/api/core/v1"

	"github.com/datawire/ambassador/pkg/kates"
	"github.com/datawire/dlib/dexec"
	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dutil"
	"github.com/datawire/telepresence2/rpc/v2/daemon"
	"github.com/datawire/telepresence2/v2/pkg/client"
)

const connectTimeout = 5 * time.Second

// k8sCluster is a Kubernetes cluster reference
type k8sCluster struct {
	// Things for bypassing kates
	Server       string
	Context      string
	Namespace    string
	kubeFlagArgs []string

	// Main
	client *kates.Client
	daemon daemon.DaemonClient // for DNS updates

	// Current snapshot.
	// These fields (except for accLock/accCond) get set by acc.Update().
	accLock     sync.Mutex
	accCond     *sync.Cond
	Deployments []*kates.Deployment
	Pods        []*kates.Pod
	Services    []*kates.Service
}

// getKubectlArgs returns the kubectl command arguments to run a
// kubectl command with this cluster.
func (kc *k8sCluster) getKubectlArgs(args ...string) []string {
	return append(kc.kubeFlagArgs, args...)
}

// getKubectlCmd returns a Cmd that runs kubectl with the given arguments and
// the appropriate environment to talk to the cluster
func (kc *k8sCluster) getKubectlCmd(c context.Context, args ...string) *dexec.Cmd {
	return dexec.CommandContext(c, "kubectl", kc.getKubectlArgs(args...)...)
}

// portForwardAndThen starts a kubectl port-forward command, passes its output to the given scanner, and waits for
// the scanner to produce a result. A new goroutine is started as a sibling in the current dgroup.Group if the scanner
// returns something other than nil. The sibling is named after the argument thenName and executes the then function,
// passing what the scanner returned as an argument.
func (kc *k8sCluster) portForwardAndThen(
	c context.Context,
	kpfArgs []string,
	outputHandler func(*bufio.Scanner) interface{},
	thenName string,
	then func(context.Context, interface{}) error,
) error {
	args := make([]string, 0, len(kc.kubeFlagArgs)+1+len(kpfArgs))
	args = append(args, kc.kubeFlagArgs...)
	args = append(args, "port-forward")
	args = append(args, kpfArgs...)

	pf := dexec.CommandContext(c, "kubectl", args...)
	out, err := pf.StdoutPipe()
	if err != nil {
		return err
	}
	defer out.Close()

	// We want this command to keep on running. If it returns an error, then it was unsuccessful.
	if err = pf.Start(); err != nil {
		out.Close()
		dlog.Errorf(c, "port-forward failed to start: %v", client.RunError(err))
		return err
	}

	sc := bufio.NewScanner(out)

	// Give port-forward 10 seconds to produce the correct output and spawn the next process
	timer := time.AfterFunc(10*time.Second, func() {
		_ = pf.Process.Kill()
	})

	// wait group is done when next process starts.
	go func() {
		defer timer.Stop()
		if output := outputHandler(sc); output != nil {
			dgroup.ParentGroup(c).Go(thenName, func(c context.Context) error { return then(c, output) })
		}
	}()

	// let the port forward continue running. It will either be killed by the
	// timer (if it didn't produce the expected output) or by a context cancel.
	if err = pf.Wait(); err != nil {
		if c.Err() != nil {
			// Context cancelled
			err = nil
		} else {
			err = client.RunError(err)
			dlog.Errorf(c, "port-forward failed: %v", err)
		}
	}
	return err
}

// check for cluster connectivity
func (kc *k8sCluster) check(c context.Context) error {
	c, cancel := context.WithTimeout(c, connectTimeout)
	defer cancel()
	cmd := kc.getKubectlCmd(c, "get", "po", "ohai", "--ignore-not-found")
	stderr := bytes.Buffer{}
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		if c.Err() == context.DeadlineExceeded {
			err = errors.New("timeout when testing cluster connectivity")
		} else if msg := stderr.String(); len(msg) > 0 {
			err = errors.New(msg)
		}
	}
	return err
}

func (kc *k8sCluster) createWatch(c context.Context, namespace string) (acc *kates.Accumulator, err error) {
	defer func() {
		if r := dutil.PanicToError(recover()); r != nil {
			err = r
		}
	}()

	return kc.client.Watch(c,
		kates.Query{
			Name:      "Services",
			Namespace: namespace,
			Kind:      "service",
		},
		kates.Query{
			Name:      "Deployments",
			Namespace: namespace,
			Kind:      "deployment",
		},
		kates.Query{
			Name:      "Pods",
			Namespace: namespace,
			Kind:      "pod",
		}), nil
}

func (kc *k8sCluster) waitUntil(ctx context.Context, fn func(context.Context) (bool, error)) error {
	kc.accLock.Lock()
	defer kc.accLock.Unlock()

	// Ensure that accCond.Wait() is released if ctx is cancelled. A child
	// context is needed since the broadcast should happen only when the
	// parent context's Err() isn't nil. The childCtx Err() will never be
	// nil due to the defer cancel() when this function returns.
	childCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	go func() {
		<-childCtx.Done()
		if ctx.Err() != nil {
			// parent context got cancelled
			kc.accCond.Broadcast()
		}
	}()

	for {
		done, err := fn(ctx)
		if done || err != nil {
			return err
		}
		kc.accCond.Wait()
		if err = ctx.Err(); err != nil {
			return err
		}
	}
}

func (kc *k8sCluster) startWatches(c context.Context, namespace string, accWait chan<- struct{}) error {
	acc, err := kc.createWatch(c, namespace)
	if err != nil {
		return err
	}

	closeAccWait := func() {
		if accWait != nil {
			close(accWait)
			accWait = nil
		}
	}

	dgroup.ParentGroup(c).Go("watch-k8s", func(c context.Context) error {
		for {
{'fraction_non_alphanumeric': 0.13207720588235294, 'fraction_numerical': 0.004963235294117647, 'mean_word_length': 3.468583162217659, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3934587', 'n_tokens_mistral': 3998, 'n_tokens_neox': 3755, 'n_words': 1111}