github.com/nsqio/nsq@v1.3.0/bench/bench_reader/bench_reader.go (about)

     1  package main
     2  
     3  import (
     4  	"bufio"
     5  	"flag"
     6  	"fmt"
     7  	"log"
     8  	"net"
     9  	"runtime"
    10  	"strings"
    11  	"sync"
    12  	"sync/atomic"
    13  	"time"
    14  
    15  	"github.com/nsqio/go-nsq"
    16  )
    17  
    18  var (
    19  	runfor     = flag.Duration("runfor", 10*time.Second, "duration of time to run")
    20  	tcpAddress = flag.String("nsqd-tcp-address", "127.0.0.1:4150", "<addr>:<port> to connect to nsqd")
    21  	size       = flag.Int("size", 200, "size of messages")
    22  	topic      = flag.String("topic", "sub_bench", "topic to receive messages on")
    23  	channel    = flag.String("channel", "ch", "channel to receive messages on")
    24  	deadline   = flag.String("deadline", "", "deadline to start the benchmark run")
    25  	rdy        = flag.Int("rdy", 2500, "RDY count to use")
    26  )
    27  
    28  var totalMsgCount int64
    29  
    30  func main() {
    31  	flag.Parse()
    32  	var wg sync.WaitGroup
    33  
    34  	log.SetPrefix("[bench_reader] ")
    35  
    36  	goChan := make(chan int)
    37  	rdyChan := make(chan int)
    38  	workers := runtime.GOMAXPROCS(0)
    39  	for j := 0; j < workers; j++ {
    40  		wg.Add(1)
    41  		go func(id int) {
    42  			subWorker(*runfor, workers, *tcpAddress, *topic, *channel, rdyChan, goChan, id)
    43  			wg.Done()
    44  		}(j)
    45  		<-rdyChan
    46  	}
    47  
    48  	if *deadline != "" {
    49  		t, err := time.Parse("2006-01-02 15:04:05", *deadline)
    50  		if err != nil {
    51  			log.Fatal(err)
    52  		}
    53  		d := time.Until(t)
    54  		log.Printf("sleeping until %s (%s)", t, d)
    55  		time.Sleep(d)
    56  	}
    57  
    58  	start := time.Now()
    59  	close(goChan)
    60  	wg.Wait()
    61  	end := time.Now()
    62  	duration := end.Sub(start)
    63  	tmc := atomic.LoadInt64(&totalMsgCount)
    64  	log.Printf("duration: %s - %.03fmb/s - %.03fops/s - %.03fus/op",
    65  		duration,
    66  		float64(tmc*int64(*size))/duration.Seconds()/1024/1024,
    67  		float64(tmc)/duration.Seconds(),
    68  		float64(duration/time.Microsecond)/float64(tmc))
    69  }
    70  
    71  func subWorker(td time.Duration, workers int, tcpAddr string, topic string, channel string, rdyChan chan int, goChan chan int, id int) {
    72  	conn, err := net.DialTimeout("tcp", tcpAddr, time.Second)
    73  	if err != nil {
    74  		panic(err.Error())
    75  	}
    76  	conn.Write(nsq.MagicV2)
    77  	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
    78  	ci := make(map[string]interface{})
    79  	ci["client_id"] = "reader"
    80  	ci["hostname"] = "reader"
    81  	ci["user_agent"] = fmt.Sprintf("bench_reader/%s", nsq.VERSION)
    82  	cmd, _ := nsq.Identify(ci)
    83  	cmd.WriteTo(rw)
    84  	nsq.Subscribe(topic, channel).WriteTo(rw)
    85  	rdyChan <- 1
    86  	<-goChan
    87  	nsq.Ready(*rdy).WriteTo(rw)
    88  	rw.Flush()
    89  	nsq.ReadResponse(rw)
    90  	nsq.ReadResponse(rw)
    91  	var msgCount int64
    92  	go func() {
    93  		time.Sleep(td)
    94  		conn.Close()
    95  	}()
    96  	for {
    97  		resp, err := nsq.ReadResponse(rw)
    98  		if err != nil {
    99  			if strings.Contains(err.Error(), "use of closed network connection") {
   100  				break
   101  			}
   102  			panic(err.Error())
   103  		}
   104  		frameType, data, err := nsq.UnpackResponse(resp)
   105  		if err != nil {
   106  			panic(err.Error())
   107  		}
   108  		if frameType == nsq.FrameTypeError {
   109  			panic(string(data))
   110  		} else if frameType == nsq.FrameTypeResponse {
   111  			continue
   112  		}
   113  		msg, err := nsq.DecodeMessage(data)
   114  		if err != nil {
   115  			panic(err.Error())
   116  		}
   117  		nsq.Finish(msg.ID).WriteTo(rw)
   118  		msgCount++
   119  		if float64(msgCount%int64(*rdy)) > float64(*rdy)*0.75 {
   120  			rw.Flush()
   121  		}
   122  	}
   123  	atomic.AddInt64(&totalMsgCount, msgCount)
   124  }