diff options
-rw-r--r-- | internal/options/options.go | 2 | ||||
-rw-r--r-- | main.go | 19 |
2 files changed, 20 insertions, 1 deletions
diff --git a/internal/options/options.go b/internal/options/options.go index d01ba83..9e1a653 100644 --- a/internal/options/options.go +++ b/internal/options/options.go @@ -21,6 +21,7 @@ type Options struct { CheckerDomain string CheckerInterval time.Duration MetricsInterval time.Duration + Limit int } func Parse() (opts Options) { @@ -36,6 +37,7 @@ func Parse() (opts Options) { flag.StringVar(&opts.CheckerDomain, "c", "", "domain with onion location for setup santity-checks") flag.DurationVar(&opts.CheckerInterval, "C", 10*time.Second, "how often to to run checker") flag.DurationVar(&opts.MetricsInterval, "m", 5*time.Second, "how often to emit metrics") + flag.IntVar(&opts.Limit, "l", 10, "rate-limit that kicks in before feeding workers in reads/s") flag.Parse() return @@ -217,6 +217,7 @@ func workGenerator(ctx context.Context, opts options.Options, fp *os.File, quest buf := make([]byte, 0, opts.MaxFileBuffer*1024*1024) scanner.Buffer(buf, opts.MaxFileBuffer*1024*1024) + // roll-up to the requested start line nextLine := int64(0) if opts.StartLineInclusive > nextLine { for scanner.Scan() { @@ -233,17 +234,32 @@ func workGenerator(ctx context.Context, opts options.Options, fp *os.File, quest } } + // initialize ticker ticker := time.NewTicker(opts.MetricsInterval) defer ticker.Stop() - startTime := time.Now().Unix() latestTime := startTime latestCount := opts.StartLineInclusive + + // initialize rate-limit + limit := time.NewTicker(time.Second) + defer limit.Stop() + readCount := 0 + for scanner.Scan() { if opts.EndLineExclusive > 0 && nextLine == opts.EndLineExclusive { break } + if readCount == opts.Limit { + select { + case <-ctx.Done(): + return nextLine, false + case <-limit.C: + readCount = 0 + } + } + select { case <-ctx.Done(): return nextLine, false @@ -259,6 +275,7 @@ func workGenerator(ctx context.Context, opts options.Options, fp *os.File, quest latestTime = now case questionCh <- qna.Question{Domain: scanner.Text()}: nextLine++ + readCount++ } } |