Add --no-follow flag to exit when all logs have been shown (#204)

Takashi Kusumi · web-flow · commit a5e581d7380a · 2023-01-14T21:27:22.000+09:00
* Extract filtering logic to targetFilter

* Add --no-follow flag to exit when all logs have been shown
diff --git a/README.md b/README.md
@@ -78,6 +78,7 @@ The `pod` query is a regular expression so you could provide `"web-\w"` to tail
  `--init-containers`         | `true`    | Include or exclude init containers.
  `--kubeconfig`              |           | Path to kubeconfig file to use. Default to KUBECONFIG variable then ~/.kube/config path.
  `--namespace`, `-n`         |           | Kubernetes namespace to use. Default to namespace configured in kubernetes context. To specify multiple namespaces, repeat this or set comma-separated value.
+ `--no-follow`               | `false`   | Exit when all logs have been shown.
  `--output`, `-o`            | `default` | Specify predefined template. Currently support: [default, raw, json, extjson, ppextjson]
  `--prompt`, `-p`            | `false`   | Toggle interactive prompt for selecting 'app.kubernetes.io/instance' label values.
  `--selector`, `-l`          |           | Selector (label query) to filter on. If present, default to ".*" for the pod-query.
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -61,6 +61,7 @@ type options struct {
 	output              string
 	prompt              bool
 	podQuery            string
+	noFollow            bool
 }
 
 func NewOptions(streams genericclioptions.IOStreams) *options {
@@ -79,6 +80,7 @@ func NewOptions(streams genericclioptions.IOStreams) *options {
 		timestamps:          false,
 		timezone:            "Local",
 		prompt:              false,
+		noFollow:            false,
 	}
 }
 
@@ -310,7 +312,7 @@ func (o *options) sternConfig() (*stern.Config, error) {
 		FieldSelector:         fieldSelector,
 		TailLines:             tailLines,
 		Template:              template,
-		Follow:                true,
+		Follow:                !o.noFollow,
 
 		Out:    o.Out,
 		ErrOut: o.ErrOut,
@@ -328,6 +330,7 @@ func (o *options) AddFlags(fs *pflag.FlagSet) {
 	fs.StringArrayVarP(&o.exclude, "exclude", "e", o.exclude, "Log lines to exclude. (regular expression)")
 	fs.StringVarP(&o.excludeContainer, "exclude-container", "E", o.excludeContainer, "Container name to exclude when multiple containers in pod. (regular expression)")
 	fs.StringVar(&o.excludePod, "exclude-pod", o.excludePod, "Pod name to exclude. (regular expression)")
+	fs.BoolVar(&o.noFollow, "no-follow", o.noFollow, "Exit when all logs have been shown.")
 	fs.StringArrayVarP(&o.include, "include", "i", o.include, "Log lines to include. (regular expression)")
 	fs.BoolVar(&o.initContainers, "init-containers", o.initContainers, "Include or exclude init containers.")
 	fs.BoolVar(&o.ephemeralContainers, "ephemeral-containers", o.ephemeralContainers, "Include or exclude ephemeral containers.")
diff --git a/go.mod b/go.mod
@@ -61,6 +61,7 @@ require (
 	golang.org/x/crypto v0.0.0-20220924013350-4ba4fb4dd9e7 // indirect
 	golang.org/x/net v0.0.0-20220923203811-8be639271d50 // indirect
 	golang.org/x/oauth2 v0.0.0-20220909003341-f21342109be1 // indirect
+	golang.org/x/sync v0.1.0 // indirect
 	golang.org/x/sys v0.0.0-20220919091848-fb04ddd9f9c8 // indirect
 	golang.org/x/term v0.0.0-20220919170432-7a66f970e087 // indirect
 	golang.org/x/text v0.3.7 // indirect
diff --git a/go.sum b/go.sum
@@ -241,6 +241,8 @@ golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
+golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
diff --git a/stern/list.go b/stern/list.go
@@ -21,6 +21,8 @@ import (
 	"github.com/pkg/errors"
 	"github.com/stern/stern/kubernetes"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/fields"
+	"k8s.io/apimachinery/pkg/labels"
 	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
 )
 
@@ -87,3 +89,18 @@ func List(ctx context.Context, config *Config) (map[string]string, error) {
 
 	return labels, nil
 }
+
+// ListTargets returns targets by listing and filtering pods
+func ListTargets(ctx context.Context, i corev1client.PodInterface, labelSelector labels.Selector, fieldSelector fields.Selector, filter *targetFilter) ([]*Target, error) {
+	list, err := i.List(ctx, metav1.ListOptions{LabelSelector: labelSelector.String(), FieldSelector: fieldSelector.String()})
+	if err != nil {
+		return nil, err
+	}
+	var targets []*Target
+	for i := range list.Items {
+		filter.visit(&list.Items[i], func(t *Target, containerStateMatched bool) {
+			targets = append(targets, t)
+		})
+	}
+	return targets, nil
+}
diff --git a/stern/stern.go b/stern/stern.go
@@ -21,6 +21,7 @@ import (
 
 	"github.com/pkg/errors"
 	"github.com/stern/stern/kubernetes"
+	"golang.org/x/sync/errgroup"
 	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
 )
 
@@ -74,6 +75,52 @@ func Run(ctx context.Context, config *Config) error {
 		}
 	}
 
+	filter := &targetFilter{
+		podFilter:              config.PodQuery,
+		excludePodFilter:       config.ExcludePodQuery,
+		containerFilter:        config.ContainerQuery,
+		containerExcludeFilter: config.ExcludeContainerQuery,
+		initContainers:         config.InitContainers,
+		ephemeralContainers:    config.EphemeralContainers,
+		containerStates:        config.ContainerStates,
+	}
+	newTail := func(t *Target) *Tail {
+		return NewTail(clientset, t.Node, t.Namespace, t.Pod, t.Container, config.Template, config.Out, config.ErrOut, &TailOptions{
+			Timestamps:   config.Timestamps,
+			Location:     config.Location,
+			SinceSeconds: int64(config.Since.Seconds()),
+			Exclude:      config.Exclude,
+			Include:      config.Include,
+			Namespace:    config.AllNamespaces || len(namespaces) > 1,
+			TailLines:    config.TailLines,
+			Follow:       config.Follow,
+		})
+	}
+
+	if !config.Follow {
+		var eg errgroup.Group
+		for _, n := range namespaces {
+			targets, err := ListTargets(ctx,
+				clientset.Pods(n),
+				config.LabelSelector,
+				config.FieldSelector,
+				filter,
+			)
+			if err != nil {
+				return err
+			}
+			for _, t := range targets {
+				t := t
+				eg.Go(func() error {
+					tail := newTail(t)
+					defer tail.Close()
+					return tail.Start(ctx)
+				})
+			}
+		}
+		return eg.Wait()
+	}
+
 	added := make(chan *Target)
 	removed := make(chan *Target)
 	errCh := make(chan error)
@@ -83,17 +130,12 @@ func Run(ctx context.Context, config *Config) error {
 	defer close(errCh)
 
 	for _, n := range namespaces {
-		a, r, err := Watch(ctx,
+		a, r, err := WatchTargets(ctx,
 			clientset.Pods(n),
-			config.PodQuery,
-			config.ExcludePodQuery,
-			config.ContainerQuery,
-			config.ExcludeContainerQuery,
-			config.InitContainers,
-			config.EphemeralContainers,
-			config.ContainerStates,
 			config.LabelSelector,
-			config.FieldSelector)
+			config.FieldSelector,
+			filter,
+		)
 		if err != nil {
 			return errors.Wrap(err, "failed to set up watch")
 		}
@@ -133,16 +175,7 @@ func Run(ctx context.Context, config *Config) error {
 				}
 			}
 
-			tail := NewTail(clientset, p.Node, p.Namespace, p.Pod, p.Container, config.Template, config.Out, config.ErrOut, &TailOptions{
-				Timestamps:   config.Timestamps,
-				Location:     config.Location,
-				SinceSeconds: int64(config.Since.Seconds()),
-				Exclude:      config.Exclude,
-				Include:      config.Include,
-				Namespace:    config.AllNamespaces || len(namespaces) > 1,
-				TailLines:    config.TailLines,
-				Follow:       config.Follow,
-			})
+			tail := newTail(p)
 			setTail(targetID, tail)
 
 			go func(tail *Tail) {
diff --git a/stern/target.go b/stern/target.go
@@ -0,0 +1,91 @@
+//   Copyright 2017 Wercker Holding BV
+//
+//   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 stern
+
+import (
+	"fmt"
+	"regexp"
+
+	corev1 "k8s.io/api/core/v1"
+)
+
+// Target is a target to watch
+type Target struct {
+	Node      string
+	Namespace string
+	Pod       string
+	Container string
+}
+
+// GetID returns the ID of the object
+func (t *Target) GetID() string {
+	return fmt.Sprintf("%s-%s-%s", t.Namespace, t.Pod, t.Container)
+}
+
+// targetFilter is a filter of Target
+type targetFilter struct {
+	podFilter              *regexp.Regexp
+	excludePodFilter       *regexp.Regexp
+	containerFilter        *regexp.Regexp
+	containerExcludeFilter *regexp.Regexp
+	initContainers         bool
+	ephemeralContainers    bool
+	containerStates        []ContainerState
+}
+
+// visit passes filtered Targets to the visitor function
+func (f *targetFilter) visit(pod *corev1.Pod, visitor func(t *Target, containerStateMatched bool)) {
+	// filter by pod
+	if !f.podFilter.MatchString(pod.Name) {
+		return
+	}
+	if f.excludePodFilter != nil && f.excludePodFilter.MatchString(pod.Name) {
+		return
+	}
+
+	// filter by container statuses
+	var statuses []corev1.ContainerStatus
+	statuses = append(statuses, pod.Status.ContainerStatuses...)
+	if f.initContainers {
+		statuses = append(statuses, pod.Status.InitContainerStatuses...)
+	}
+	if f.ephemeralContainers {
+		statuses = append(statuses, pod.Status.EphemeralContainerStatuses...)
+	}
+	for _, c := range statuses {
+		if !f.containerFilter.MatchString(c.Name) {
+			continue
+		}
+		if f.containerExcludeFilter != nil && f.containerExcludeFilter.MatchString(c.Name) {
+			continue
+		}
+		t := &Target{
+			Node:      pod.Spec.NodeName,
+			Namespace: pod.Namespace,
+			Pod:       pod.Name,
+			Container: c.Name,
+		}
+		visitor(t, f.matchContainerState(c.State))
+	}
+}
+
+func (f *targetFilter) matchContainerState(state corev1.ContainerState) bool {
+	for _, containerState := range f.containerStates {
+		if containerState.Match(state) {
+			return true
+		}
+	}
+	return false
+}
diff --git a/stern/target_test.go b/stern/target_test.go
diff --git a/stern/watch.go b/stern/watch.go
