Implements feature to stream pod logs in Loki

cmd/run.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/pkg/errors"
@@ -16,24 +17,28 @@ import (
 )
 
 type runOptions struct {
-	kubeconfig string
+	kubeconfig     string
+	streamDuration time.Duration
 }
 
 var runOpts = &runOptions{}
 
 var runCmd = &cobra.Command{
-	Use:     "run",
-	Short:   "run ",
-	Long:    "run ",
-	Example: "	lomba run",
-	//Args: cobra.MaximumNArgs(1),
+	Use:   "run",
+	Short: "run ",
+	Long:  "run ",
+	Example: "	lomba run\n" +
+		"	lomba run --kubeconfig <path to kubeconfig> --stream-duration 10m\n" +
+		"	lomba run -k <path to kubeconfig> -s 15m30s",
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return runRun()
 	},
 }
 
 func init() {
-	runCmd.Flags().StringVarP(&runOpts.kubeconfig, "kubeconfig", "k", "", "kubeconfig")
+	runCmd.Flags().StringVarP(&runOpts.kubeconfig, "kubeconfig", "k", os.Getenv("HOME")+"/.kube/config", "kubeconfig")
+	runCmd.Flags().DurationVarP(&runOpts.streamDuration, "stream-duration", "s", 1*time.Hour, "time duration to "+
+		"stream logs into loki")
 	RootCmd.AddCommand(runCmd)
 }
 
@@ -53,15 +58,21 @@ func runRun() error {
 		return err
 	}
 
-	err = rr.Run(context.Background())
+	cancelCtx, cancelFunc := context.WithTimeout(context.Background(), runOpts.streamDuration)
+	err = rr.Run(cancelCtx)
 	if err != nil {
+		cancelFunc()
 		return err
 	}
 
 	grafanaEndpoint := grafana.GetOutboundIPOrLocalhost()
 
-	fmt.Printf("Kubernetes logs are injested to Loki. Ready to query at http://%s:3000\n", grafanaEndpoint)
+	fmt.Printf("Kubernetes logs will be streamed to Loki for next %s minutes. Ready to query at http://%s:3000\n",
+		runOpts.streamDuration, grafanaEndpoint)
 
+	// sleep for duration as much as set in stream-duration flag to keep the goroutines active
+	time.Sleep(runOpts.streamDuration)
+	cancelFunc()
 	return nil
 }
 

go.mod

@@ -10,6 +10,8 @@ require (
 	github.com/prometheus/client_golang v1.12.1
 	github.com/prometheus/common v0.32.1
 	github.com/spf13/cobra v1.0.0
+	k8s.io/api v0.23.6
+	k8s.io/apimachinery v0.23.6
 	k8s.io/client-go v0.23.6
 )
 
@@ -125,8 +127,6 @@ require (
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
-	k8s.io/api v0.23.6 // indirect
-	k8s.io/apimachinery v0.23.6 // indirect
 	k8s.io/klog/v2 v2.40.1 // indirect
 	k8s.io/kube-openapi v0.0.0-20211115234752-e816edb12b65 // indirect
 	k8s.io/utils v0.0.0-20211116205334-6203023598ed // indirect

internal/runner/runner.go

@@ -2,9 +2,7 @@ package runner
 
 import (
 	"bufio"
-	"bytes"
 	"context"
-	"io"
 
 	"github.com/go-kit/log"
 	corev1 "k8s.io/api/core/v1"
@@ -48,8 +46,7 @@ func NewRunner(params RunnerParams) (Runner, error) {
 		cs:         params.ClientSet}, nil
 }
 
-func (r *runner) Run(ctx context.Context) error {
-
+func (r *runner) Run(cancelCtx context.Context) error {
 	err := loki.BringUpPod()
 	if err != nil {
 		return err
@@ -60,68 +57,61 @@ func (r *runner) Run(ctx context.Context) error {
 		return err
 	}
 
-	namespaceList, err := r.cs.CoreV1().Namespaces().List(ctx, metav1.ListOptions{})
+	// get list of pods from all namespaces
+	podList, err := r.cs.CoreV1().Pods("").List(cancelCtx, metav1.ListOptions{})
 	if err != nil {
 		return err
 	}
 
-	for _, ns := range namespaceList.Items {
-		podList, err := r.cs.CoreV1().Pods(ns.Name).List(ctx, metav1.ListOptions{})
-		if err != nil {
-			return err
-		}
+	for _, pod := range podList.Items {
+		go r.streamPodLogs(cancelCtx, r.cs, pod)
 
-		for _, pod := range podList.Items {
-
-			for _, container := range pod.Spec.Containers {
-				req := r.cs.CoreV1().Pods(ns.Name).GetLogs(pod.Name, &corev1.PodLogOptions{
-					Timestamps: true,
-					Container:  container.Name,
-				})
-
-				podLogs, err := req.Stream(ctx)
-				if err != nil {
-					return err
-				}
-				defer podLogs.Close()
-
-				buf := new(bytes.Buffer)
-
-				_, err = io.Copy(buf, podLogs)
-				if err != nil {
-					return err
-				}
-
-				labels := make(map[string]string)
-				labels["namespace"] = ns.Name
-				labels["pod_name"] = pod.Name
-				labels["container_name"] = container.Name
-
-				err = r.loadLogsToLoki(buf, parser.NewContainerParser(), labels)
-				if err != nil {
-					return err
-				}
-			}
-		}
 	}
 
 	return nil
 }
 
-func (r *runner) loadLogsToLoki(rawLogs *bytes.Buffer, logParser parser.Parser, labels map[string]string) error {
+func (r *runner) loadLogsToLoki(logLine string, logParser parser.Parser, labels map[string]string) error {
+	tm, labelset, err := logParser.Parse(logLine, labels)
+	if err != nil {
+		r.logger.Log("Skipping log due to invalid parse", "Error", err.Error())
+		return err
+	}
+	r.lokiClient.PostLog(logLine, tm, labelset)
 
-	scanner := bufio.NewScanner(rawLogs)
-	scanner.Split(bufio.ScanLines)
+	return nil
+}
 
-	for scanner.Scan() {
-		log_line := scanner.Text()
-		tm, labels, err := logParser.Parse(log_line, labels)
+// streamPodLogs will stream the pod logs and load the logs to loki with relevant
+// labels, loglines and timestamp
+func (r *runner) streamPodLogs(cancelCtx context.Context, cs kubernetes.Interface, pod corev1.Pod) error {
+	for _, container := range pod.Spec.Containers {
+		podLogOptions := &corev1.PodLogOptions{
+			Follow:     true,
+			Timestamps: true,
+		}
+
+		req := cs.CoreV1().Pods(pod.Namespace).GetLogs(pod.Name, podLogOptions)
+		stream, err := req.Stream(cancelCtx)
 		if err != nil {
-			r.logger.Log("Skipping log due to invalid parse", "Error", err.Error())
-			continue
+			return err
 		}
-		r.lokiClient.PostLog(log_line, tm, labels)
-	}
 
+		reader := bufio.NewScanner(stream)
+		reader.Split(bufio.ScanLines)
+		defer stream.Close()
+
+		for reader.Scan() {
+			labels := make(map[string]string)
+			labels["namespace"] = pod.Namespace
+			labels["pod_name"] = pod.Name
+			labels["container_name"] = container.Name
+
+			logLine := reader.Text()
+
+			// ignore the error and continue reading stream & loading to loki
+			_ = r.loadLogsToLoki(logLine, parser.NewContainerParser(), labels)
+		}
+	}
 	return nil
 }