github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/clients/pkg/promtail/targets/stdin/stdin_target_manager.go (about)

     1  package stdin
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"fmt"
     7  	"io"
     8  	"os"
     9  	"strings"
    10  	"time"
    11  
    12  	"github.com/go-kit/log"
    13  	"github.com/go-kit/log/level"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/prometheus/common/model"
    16  	"github.com/prometheus/prometheus/discovery/targetgroup"
    17  
    18  	"github.com/grafana/loki/clients/pkg/logentry/stages"
    19  	"github.com/grafana/loki/clients/pkg/promtail/api"
    20  	"github.com/grafana/loki/clients/pkg/promtail/scrapeconfig"
    21  	"github.com/grafana/loki/clients/pkg/promtail/targets/target"
    22  
    23  	"github.com/grafana/loki/pkg/logproto"
    24  )
    25  
    26  // bufferSize is the size of the buffered reader
    27  const bufferSize = 8096
    28  
    29  // file is an interface allowing us to abstract a file.
    30  type file interface {
    31  	Stat() (os.FileInfo, error)
    32  	io.Reader
    33  }
    34  
    35  var (
    36  	// stdIn is os.Stdin but can be replaced for testing purpose.
    37  	stdIn       file = os.Stdin
    38  	hostName, _      = os.Hostname()
    39  	// defaultStdInCfg is the default config for stdin target if none provided.
    40  	defaultStdInCfg = scrapeconfig.Config{
    41  		JobName: "stdin",
    42  		ServiceDiscoveryConfig: scrapeconfig.ServiceDiscoveryConfig{
    43  			StaticConfigs: []*targetgroup.Group{
    44  				{Labels: model.LabelSet{"job": "stdin"}},
    45  				{Labels: model.LabelSet{"hostname": model.LabelValue(hostName)}},
    46  			},
    47  		},
    48  	}
    49  )
    50  
    51  type Shutdownable interface {
    52  	Shutdown()
    53  }
    54  
    55  // nolint:revive
    56  type StdinTargetManager struct {
    57  	*readerTarget
    58  	app Shutdownable
    59  }
    60  
    61  func NewStdinTargetManager(reg prometheus.Registerer, log log.Logger, app Shutdownable, client api.EntryHandler, configs []scrapeconfig.Config) (*StdinTargetManager, error) {
    62  	reader, err := newReaderTarget(reg, log, stdIn, client, getStdinConfig(log, configs))
    63  	if err != nil {
    64  		return nil, err
    65  	}
    66  	stdinManager := &StdinTargetManager{
    67  		readerTarget: reader,
    68  		app:          app,
    69  	}
    70  	// when we're done flushing our stdin we can shutdown the app.
    71  	go func() {
    72  		<-reader.ctx.Done()
    73  		app.Shutdown()
    74  	}()
    75  	return stdinManager, nil
    76  }
    77  
    78  func getStdinConfig(log log.Logger, configs []scrapeconfig.Config) scrapeconfig.Config {
    79  	cfg := defaultStdInCfg
    80  	// if we receive configs we use the first one.
    81  	if len(configs) > 0 {
    82  		if len(configs) > 1 {
    83  			level.Warn(log).Log("msg", fmt.Sprintf("too many scrape configs, skipping %d configs.", len(configs)-1))
    84  		}
    85  		cfg = configs[0]
    86  	}
    87  	return cfg
    88  }
    89  
    90  func (t *StdinTargetManager) Ready() bool {
    91  	return t.ctx.Err() == nil
    92  }
    93  func (t *StdinTargetManager) Stop()                                     { t.cancel() }
    94  func (t *StdinTargetManager) ActiveTargets() map[string][]target.Target { return nil }
    95  func (t *StdinTargetManager) AllTargets() map[string][]target.Target    { return nil }
    96  
    97  type readerTarget struct {
    98  	in     *bufio.Reader
    99  	out    api.EntryHandler
   100  	lbs    model.LabelSet
   101  	logger log.Logger
   102  
   103  	cancel context.CancelFunc
   104  	ctx    context.Context
   105  }
   106  
   107  func newReaderTarget(reg prometheus.Registerer, logger log.Logger, in io.Reader, client api.EntryHandler, cfg scrapeconfig.Config) (*readerTarget, error) {
   108  	pipeline, err := stages.NewPipeline(log.With(logger, "component", "pipeline"), cfg.PipelineStages, &cfg.JobName, reg)
   109  	if err != nil {
   110  		return nil, err
   111  	}
   112  	lbs := model.LabelSet{}
   113  	for _, static := range cfg.ServiceDiscoveryConfig.StaticConfigs {
   114  		if static != nil && static.Labels != nil {
   115  			lbs = lbs.Merge(static.Labels)
   116  		}
   117  	}
   118  	ctx, cancel := context.WithCancel(context.Background())
   119  	t := &readerTarget{
   120  		in:     bufio.NewReaderSize(in, bufferSize),
   121  		out:    pipeline.Wrap(client),
   122  		cancel: cancel,
   123  		ctx:    ctx,
   124  		lbs:    lbs,
   125  		logger: log.With(logger, "component", "reader"),
   126  	}
   127  	go t.read()
   128  
   129  	return t, nil
   130  }
   131  
   132  func (t *readerTarget) read() {
   133  	defer t.cancel()
   134  	defer t.out.Stop()
   135  
   136  	entries := t.out.Chan()
   137  	for {
   138  		if t.ctx.Err() != nil {
   139  			return
   140  		}
   141  		line, err := t.in.ReadString('\n')
   142  		if err != nil && err != io.EOF {
   143  			level.Warn(t.logger).Log("msg", "error reading buffer", "err", err)
   144  			return
   145  		}
   146  		line = strings.TrimRight(line, "\r\n")
   147  		if line == "" {
   148  			if err == io.EOF {
   149  				return
   150  			}
   151  			continue
   152  		}
   153  		entries <- api.Entry{
   154  			Labels: t.lbs.Clone(),
   155  			Entry: logproto.Entry{
   156  				Timestamp: time.Now(),
   157  				Line:      line,
   158  			},
   159  		}
   160  		if err == io.EOF {
   161  			return
   162  		}
   163  	}
   164  }