github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/operator/builtin/input/generate.go (about)

     1  package input
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/observiq/carbon/entry"
    10  	"github.com/observiq/carbon/operator"
    11  	"github.com/observiq/carbon/operator/helper"
    12  )
    13  
    14  func init() {
    15  	operator.Register("generate_input", func() operator.Builder { return NewGenerateInputConfig("") })
    16  }
    17  
    18  func NewGenerateInputConfig(operatorID string) *GenerateInputConfig {
    19  	return &GenerateInputConfig{
    20  		InputConfig: helper.NewInputConfig(operatorID, "generate_input"),
    21  	}
    22  }
    23  
    24  // GenerateInputConfig is the configuration of a generate input operator.
    25  type GenerateInputConfig struct {
    26  	helper.InputConfig `yaml:",inline"`
    27  	Entry              entry.Entry `json:"entry"           yaml:"entry"`
    28  	Count              int         `json:"count,omitempty" yaml:"count,omitempty"`
    29  	Static             bool        `json:"static"          yaml:"static,omitempty"`
    30  }
    31  
    32  // Build will build a generate input operator.
    33  func (c *GenerateInputConfig) Build(context operator.BuildContext) (operator.Operator, error) {
    34  	inputOperator, err := c.InputConfig.Build(context)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  
    39  	c.Entry.Record = recursiveMapInterfaceToMapString(c.Entry.Record)
    40  
    41  	generateInput := &GenerateInput{
    42  		InputOperator: inputOperator,
    43  		entry:         c.Entry,
    44  		count:         c.Count,
    45  		static:        c.Static,
    46  	}
    47  	return generateInput, nil
    48  }
    49  
    50  // GenerateInput is an operator that generates log entries.
    51  type GenerateInput struct {
    52  	helper.InputOperator
    53  	entry  entry.Entry
    54  	count  int
    55  	static bool
    56  	cancel context.CancelFunc
    57  	wg     *sync.WaitGroup
    58  }
    59  
    60  // Start will start generating log entries.
    61  func (g *GenerateInput) Start() error {
    62  	ctx, cancel := context.WithCancel(context.Background())
    63  	g.cancel = cancel
    64  	g.wg = &sync.WaitGroup{}
    65  
    66  	g.wg.Add(1)
    67  	go func() {
    68  		defer g.wg.Done()
    69  		i := 0
    70  		for {
    71  			select {
    72  			case <-ctx.Done():
    73  				return
    74  			default:
    75  			}
    76  
    77  			entry := g.entry.Copy()
    78  			if !g.static {
    79  				entry.Timestamp = time.Now()
    80  			}
    81  			g.Write(ctx, entry)
    82  
    83  			i++
    84  			if i == g.count {
    85  				return
    86  			}
    87  		}
    88  	}()
    89  
    90  	return nil
    91  }
    92  
    93  // Stop will stop generating logs.
    94  func (g *GenerateInput) Stop() error {
    95  	g.cancel()
    96  	g.wg.Wait()
    97  	return nil
    98  }
    99  
   100  func recursiveMapInterfaceToMapString(m interface{}) interface{} {
   101  	switch m := m.(type) {
   102  	case map[string]interface{}:
   103  		newMap := make(map[string]interface{})
   104  		for k, v := range m {
   105  			newMap[k] = recursiveMapInterfaceToMapString(v)
   106  		}
   107  		return newMap
   108  	case map[interface{}]interface{}:
   109  		newMap := make(map[string]interface{})
   110  		for k, v := range m {
   111  			kStr, ok := k.(string)
   112  			if !ok {
   113  				kStr = fmt.Sprintf("%v", k)
   114  			}
   115  			newMap[kStr] = recursiveMapInterfaceToMapString(v)
   116  		}
   117  		return newMap
   118  	default:
   119  		return m
   120  	}
   121  }