@@ -5,8 +5,11 @@ import (
55 "fmt"
66 "math/rand"
77 "os"
8+ "sync"
89 "time"
910
11+ "github.com/prometheus/client_golang/prometheus"
12+
1013 "github.com/ViaQ/cluster-logging-load-client/internal/clients"
1114
1215 "github.com/elastic/go-elasticsearch/v6"
@@ -29,7 +32,7 @@ const (
2932 ElasticsearchClientType ClientType = "elasticsearch"
3033)
3134
32- // Options describes the settings that can be modified for the querier
35+ // Options describes the settings that can be modified for the log generator
3336type Options struct {
3437 // Client describes the client to use for forwarding
3538 Client ClientType
@@ -43,6 +46,12 @@ type Options struct {
4346 DisableSecurityCheck bool
4447 // LogsPerSecond is the number of logs to write per second
4548 LogsPerSecond int
49+
50+ LogType string
51+ LogFormat string
52+ LabelType string
53+ SyntheticPayloadSize int
54+ UseRandomHostname bool
4655}
4756
4857// LogGenerator describes an object which generates logs
@@ -54,12 +63,22 @@ type LogGenerator struct {
5463 rate int
5564 writeToDestination func (string , string , LabelSetOptions ) error
5665 deferClose func ()
66+ logCount prometheus.Counter
67+ opts Options
5768}
5869
59- func NewLogGenerator (opts Options ) (* LogGenerator , error ) {
70+ func NewLogGenerator (opts Options , registry * prometheus. Registry ) (* LogGenerator , error ) {
6071 generator := LogGenerator {
72+ opts : opts ,
6173 rate : opts .LogsPerSecond ,
74+ logCount : prometheus .NewCounter (prometheus.CounterOpts {
75+ Name : "log_generator_messages_produced_total" ,
76+ Help : "Total number of messages produced by the log generator" ,
77+ }),
6278 }
79+ registry .MustRegister (
80+ generator .logCount ,
81+ )
6382
6483 switch opts .Client {
6584 case "file" :
@@ -113,40 +132,53 @@ func NewLogGenerator(opts Options) (*LogGenerator, error) {
113132 return & generator , nil
114133}
115134
116- func (g * LogGenerator ) GenerateLogs (logType LogType , logFormat Format , logSize int , labelOpts LabelSetOptions , randomizeHostname bool ) {
135+ func (g * LogGenerator ) Start (ctx context.Context , wg * sync.WaitGroup , errCh chan <- error ) {
136+ wg .Add (1 )
137+ go func () {
138+ <- ctx .Done ()
139+ log .Debug ("Shutting down log generator..." )
140+ g .deferClose ()
141+ }()
142+ go func () {
143+ defer wg .Done ()
144+ g .GenerateLogs ()
145+ }()
146+ }
147+
148+ func (g * LogGenerator ) GenerateLogs () {
117149 host , err := os .Hostname ()
118150 if err != nil {
119151 log .Fatalf ("error getting hostname: %s" , err )
120152 }
121-
122153 defer g .deferClose ()
123154
124155 var lineCount int64 = 0
125156
126157 logHostname := host
127- if randomizeHostname {
158+ if g . opts . UseRandomHostname {
128159 logHostname = fmt .Sprintf ("%s.%032X" , host , rand .Uint64 ())
129160 }
130161
131162 for {
132163 next := time .Now ().UTC ().Add (1 * time .Second )
133164
134165 for i := 0 ; i < g .rate ; i ++ {
135- logLine , err := RandomLog (logType , logSize )
166+ logLine , err := RandomLog (LogType ( g . opts . LogType ), g . opts . SyntheticPayloadSize )
136167 if err != nil {
137168 log .Fatalf ("error creating log: %s" , err )
138169 }
139170
140- formattedLogLine , err := FormatLog (logFormat , logHostname , lineCount , logLine )
171+ formattedLogLine , err := FormatLog (Format ( g . opts . LogFormat ) , logHostname , lineCount , logLine )
141172 if err != nil {
142173 log .Fatalf ("error formating log: %s" , err )
143174 }
144175
145- err = g .writeToDestination (host , formattedLogLine , labelOpts )
176+ err = g .writeToDestination (host , formattedLogLine , LabelSetOptions ( g . opts . LabelType ) )
146177 if err != nil {
147178 log .Fatalf ("error writing log: %s" , err )
148179 }
149180
181+ g .logCount .Inc ()
150182 lineCount ++
151183 }
152184
0 commit comments