流量回放工具之GoReplay input_file 源码分析

  -input-file value   //从一个文件中读取请求
        Read requests from file: 
                gor --input-file ./requests.gor --output-http staging.com
  -input-file-dry-run  //模拟从数据源读取数据而不重新回放它
        Simulate reading from the data source without replaying it. You will get information about expected replay time, number of found records etc.
  -input-file-loop  //循环读取文件
        Loop input files, useful for performance testing.
  -input-file-max-wait duration  //设置请求之间的最大时间间隔
        Set the maximum time between requests. Can help in situations when you have too long periods between request, and you want to skip them. Example: --input-raw-max-wait 1s
  -input-file-read-depth int  //尝试提前读取和缓存多个记录。与此同时，如果请求没有按顺序出现，它还可以对请求进行排序。因为它需要在内存中保存这个缓冲区，较大的值会导致更差的性能(默认为100)
        GoReplay tries to read and cache multiple records, in advance. In parallel it also perform sorting of requests, if they came out of order. Since it needs hold this buffer in memory, bigger values can cause worse performance (default 100)

[root@vm-1 ~]./gor --input-file "requests.gor|200%" --output-http="http://172.16.106.237:8082"
2021/08/17 15:03:58 [PPID 12356 and PID 18187] Version:1.3.0
[DEBUG][elapsed 1.361742ms]: [INPUT-FILE] FileInput: end of file 'requests.gor'

func main() {
 //获取环境变量
 if os.Getenv("GOMAXPROCS") == "" {
  //设置最大逻辑核
  runtime.GOMAXPROCS(runtime.NumCPU() * 2)
 }

 args := os.Args[1:]
 var plugins *InOutPlugins
 //文件服务
 if len(args) > 0 && args[0] == "file-server" {
  if len(args) != 2 {
   log.Fatal("You should specify port and IP (optional) for the file server. Example: `gor file-server :80`")
  }
  dir, _ := os.Getwd()

  Debug(0, "Started example file server for current directory on address ", args[1])

  log.Fatal(http.ListenAndServe(args[1], loggingMiddleware(args[1], http.FileServer(http.Dir(dir)))))
 } else {
  //解析命令行参数
  flag.Parse()
  //初始化全局的 Settings 变量
  checkSettings()
  //初始化可用插件
  plugins = NewPlugins()
 }

 log.Printf("[PPID %d and PID %d] Version:%s
", os.Getppid(), os.Getpid(), VERSION)

 if len(plugins.Inputs) == 0 || len(plugins.Outputs) == 0 {
  log.Fatal("Required at least 1 input and 1 output")
 }

 if *memprofile != "" {
  profileMEM(*memprofile)
 }

 if *cpuprofile != "" {
  profileCPU(*cpuprofile)
 }

 if Settings.Pprof != "" {
  go func() {
   log.Println(http.ListenAndServe(Settings.Pprof, nil))
  }()
 }

 closeCh := make(chan int)
 //程序核心事件处理
 emitter := NewEmitter()
 //调用 Start 函数，启动 emitter
 go emitter.Start(plugins, Settings.Middleware)
 if Settings.ExitAfter > 0 {
  log.Printf("Running gor for a duration of %s
", Settings.ExitAfter)

  time.AfterFunc(Settings.ExitAfter, func() {
   log.Printf("gor run timeout %s
", Settings.ExitAfter)
   close(closeCh)
  })
 }
 c := make(chan os.Signal, 1)
 signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 exit := 0
 select {
 case <-c:
  exit = 1
 case <-closeCh:
  exit = 0
 }
 //关闭所有协程
 emitter.Close()
 os.Exit(exit)
}

// Automatically detects type of plugin and initialize it
//
// See this article if curious about reflect stuff below: http://blog.burntsushi.net/type-parametric-functions-golang
func (plugins *InOutPlugins) registerPlugin(constructor interface{}, options ...interface{}) {
 var path, limit string
 vc := reflect.ValueOf(constructor)

 // Pre-processing options to make it work with reflect
 vo := []reflect.Value{}
 for _, oi := range options {
  vo = append(vo, reflect.ValueOf(oi))
 }

 if len(vo) > 0 {
  // Removing limit options from path
  path, limit = extractLimitOptions(vo[0].String())

  // Writing value back without limiter "|" options
  vo[0] = reflect.ValueOf(path)
 }

 // Calling our constructor with list of given options
 plugin := vc.Call(vo)[0].Interface()

 if limit != "" {
  plugin = NewLimiter(plugin, limit)
 }

 // Some of the output can be Readers as well because return responses
 if r, ok := plugin.(PluginReader); ok {
  plugins.Inputs = append(plugins.Inputs, r)
 }

 if w, ok := plugin.(PluginWriter); ok {
  plugins.Outputs = append(plugins.Outputs, w)
 }
 plugins.All = append(plugins.All, plugin)
}

// extractLimitOptions detects if plugin get called with limiter support
// Returns address and limit
func extractLimitOptions(options string) (string, string) {
 split := strings.Split(options, "|")

 if len(split) > 1 {
  return split[0], split[1]
 }

 return split[0], ""
}

package main

import (
 "fmt"
 "io"
 "math/rand"
 "strconv"
 "strings"
 "time"
)

// Limiter is a wrapper for input or output plugin which adds rate limiting
type Limiter struct {
 plugin    interface{}
 limit     int
 isPercent bool

 currentRPS  int
 currentTime int64
}

func parseLimitOptions(options string) (limit int, isPercent bool) {
 if n := strings.Index(options, "%"); n > 0 {
  limit, _ = strconv.Atoi(options[:n])
  isPercent = true
 } else {
  limit, _ = strconv.Atoi(options)
  isPercent = false
 }

 return
}

// NewLimiter constructor for Limiter, accepts plugin and options
// `options` allow to sprcify relatve or absolute limiting
func NewLimiter(plugin interface{}, options string) PluginReadWriter {
 l := new(Limiter)
 l.limit, l.isPercent = parseLimitOptions(options)
 l.plugin = plugin
 l.currentTime = time.Now().UnixNano()

 // FileInput have its own rate limiting. Unlike other inputs we not just dropping requests, we can slow down or speed up request emittion.
 if fi, ok := l.plugin.(*FileInput); ok && l.isPercent {
  fi.speedFactor = float64(l.limit) / float64(100)
 }

 return l
}

func (l *Limiter) isLimited() bool {
 // File input have its own limiting algorithm
 if _, ok := l.plugin.(*FileInput); ok && l.isPercent {
  return false
 }

 if l.isPercent {
  return l.limit <= rand.Intn(100)
 }

 if (time.Now().UnixNano() - l.currentTime) > time.Second.Nanoseconds() {
  l.currentTime = time.Now().UnixNano()
  l.currentRPS = 0
 }

 if l.currentRPS >= l.limit {
  return true
 }

 l.currentRPS++

 return false
}

// PluginWrite writes message to this plugin
func (l *Limiter) PluginWrite(msg *Message) (n int, err error) {
 if l.isLimited() {
  return 0, nil
 }
 if w, ok := l.plugin.(PluginWriter); ok {
  return w.PluginWrite(msg)
 }
 // avoid further writing
 return 0, io.ErrClosedPipe
}

// PluginRead reads message from this plugin
func (l *Limiter) PluginRead() (msg *Message, err error) {
 if r, ok := l.plugin.(PluginReader); ok {
  msg, err = r.PluginRead()
 } else {
  // avoid further reading
  return nil, io.ErrClosedPipe
 }

 if l.isLimited() {
  return nil, nil
 }

 return
}

func (l *Limiter) String() string {
 return fmt.Sprintf("Limiting %s to: %d (isPercent: %v)", l.plugin, l.limit, l.isPercent)
}

// Close closes the resources.
func (l *Limiter) Close() error {
 if fi, ok := l.plugin.(io.Closer); ok {
  fi.Close()
 }
 return nil
}

$ sudo ./gor --input-file 'request.gor|10000%' --input-file-loop --output-http 'http://10.96.136.36:8201' 

 for _, options := range Settings.InputFile {
  plugins.registerPlugin(NewFileInput, options, Settings.InputFileLoop, Settings.InputFileReadDepth, Settings.InputFileMaxWait, Settings.InputFileDryRun)
 }

// NewFileInput constructor for FileInput. Accepts file path as argument.
func NewFileInput(path string, loop bool, readDepth int, maxWait time.Duration, dryRun bool) (i *FileInput) {
 i = new(FileInput)
 i.data = make(chan []byte, 1000)
 i.exit = make(chan bool)
 i.path = path
 i.speedFactor = 1
 i.loop = loop
 i.readDepth = readDepth
 i.stats = expvar.NewMap("file-" + path)
 i.dryRun = dryRun
 i.maxWait = maxWait

 if err := i.init(); err != nil {
  return
 }

 go i.emit()

 return
}

// FileInput can read requests generated by FileOutput
type FileInput struct {
 mu          sync.Mutex
 data        chan []byte
 exit        chan bool
 path        string
 readers     []*fileInputReader
 speedFactor float64
 loop        bool
 readDepth   int
 dryRun      bool
 maxWait     time.Duration

 stats *expvar.Map
}

  if reader == nil {
   if i.loop {
    i.init()
    lastTime = -1
    continue
   } else {
    break
   }
  }