[add] improved subscriber performance. displaying received message rate on subscriber

Author: filipecosta90

.gitignore

@@ -82,3 +82,8 @@ bin/
 .Trashes
 ehthumbs.db
 Thumbs.db
+
+
+# Json Results #
+################
+*.json

subscriber.go

@@ -1,48 +1,54 @@
 package main
 
 import (
+	"encoding/json"
 	"flag"
 	"fmt"
 	"github.com/mediocregopher/radix"
+	"io/ioutil"
 	"log"
 	"math/rand"
 	"os"
 	"os/signal"
-	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"text/tabwriter"
 	"time"
 )
 
-func subscriberRoutine(addr string, subscriberName string, channel string, printMessages bool, stop chan struct{}, wg *sync.WaitGroup) {
+var totalMessages uint64
+
+type testResult struct {
+	StartTime             int64     `json:"StartTime"`
+	Duration              float64   `json:"Duration"`
+	MessageRate           float64   `json:"MessageRate"`
+	TotalMessages         uint64    `json:"TotalMessages"`
+	TotalSubscriptions    int       `json:"TotalSubscriptions"`
+	ChannelMin            int       `json:"ChannelMin"`
+	ChannelMax            int       `json:"ChannelMax"`
+	SubscribersPerChannel int       `json:"SubscribersPerChannel"`
+	MessagesPerChannel    int64     `json:"MessagesPerChannel"`
+	MessageRateTs         []float64 `json:"MessageRateTs"`
+	OSSDistributedSlots   bool      `json:"OSSDistributedSlots"`
+	Addresses             []string  `json:"Addresses"`
+}
+
+func subscriberRoutine(addr string, subscriberName string, channel string, messages_per_channel int64, printMessages bool, stop chan struct{}, wg *sync.WaitGroup) {
 	// tell the caller we've stopped
 	defer wg.Done()
 
-	conn, err, ps, msgCh, tick := bootstrapPubSub(addr, subscriberName, channel)
+	conn, _, _, msgCh, tick := bootstrapPubSub(addr, subscriberName, channel)
 	defer conn.Close()
 	defer tick.Stop()
 
 	for {
 		select {
-		case <-tick.C:
-			err = ps.Ping()
-			if err != nil {
-				//try to bootstrap again
-				conn, err, ps, msgCh, tick = bootstrapPubSub(addr, subscriberName, channel)
-				defer conn.Close()
-				defer tick.Stop()
-				if err != nil {
-					panic(err)
-				}
-			}
-			// loop
-
 		case msg := <-msgCh:
 			if printMessages {
 				fmt.Println(fmt.Sprintf("received message in channel %s. Message: %s", msg.Channel, msg.Message))
 			}
-
+			atomic.AddUint64(&totalMessages, 1)
 			break
 		case <-stop:
 			return
@@ -70,7 +76,7 @@ func bootstrapPubSub(addr string, subscriberName string, channel string) (radix.
 	if err != nil {
 		panic(err)
 	}
-	tickTime := 10 + rand.Intn(10)
+	tickTime := 1 + rand.Intn(2)
 	tick := time.NewTicker(time.Duration(tickTime) * time.Second)
 
 	return conn, err, ps, msgCh, tick
@@ -82,16 +88,24 @@ func main() {
 	channel_minimum := flag.Int("channel-minimum", 1, "channel ID minimum value ( each channel has a dedicated thread ).")
 	channel_maximum := flag.Int("channel-maximum", 100, "channel ID maximum value ( each channel has a dedicated thread ).")
 	subscribers_per_channel := flag.Int("subscribers-per-channel", 1, "number of subscribers per channel.")
+	messages_per_channel_subscriber := flag.Int64("messages", 0, "Number of total messages per subscriber per channel.")
+	json_out_file := flag.String("json-out-file", "", "Name of json output file, if not set, will not print to json.")
 	client_update_tick := flag.Int("client-update-tick", 1, "client update tick.")
+	test_time := flag.Int("test-time", 0, "Number of seconds to run the test, after receiving the first message.")
 	subscribe_prefix := flag.String("subscriber-prefix", "channel-", "prefix for subscribing to channel, used in conjunction with key-minimum and key-maximum.")
 	client_output_buffer_limit_pubsub := flag.String("client-output-buffer-limit-pubsub", "", "Specify client output buffer limits for clients subscribed to at least one pubsub channel or pattern. If the value specified is different that the one present on the DB, this setting will apply.")
 	distributeSubscribers := flag.Bool("oss-cluster-api-distribute-subscribers", false, "read cluster slots and distribute subscribers among them.")
 	printMessages := flag.Bool("print-messages", false, "print messages.")
 	flag.Parse()
-
+	totalMessages = 0
 	var nodes []radix.ClusterNode
+	var nodesAddresses []string
 	var node_subscriptions_count []int
 
+	if *test_time != 0 && *messages_per_channel_subscriber != 0 {
+		log.Fatal(fmt.Errorf("--messages and --test-time are mutially exclusive ( please specify one or the other )"))
+	}
+
 	if *distributeSubscribers {
 		// Create a normal redis connection
 		conn, err := radix.Dial("tcp", fmt.Sprintf("%s:%s", *host, *port))
@@ -111,6 +125,7 @@ func main() {
 				slot.Addr = fmt.Sprintf("%s:%s", *host, slot_port)
 			}
 			nodes = append(nodes, slot)
+			nodesAddresses = append(nodesAddresses, slot.Addr)
 			node_subscriptions_count = append(node_subscriptions_count, 0)
 		}
 		conn.Close()
@@ -126,6 +141,7 @@ func main() {
 				SecondaryOfID:   "",
 			}
 			nodes = append(nodes, node)
+			nodesAddresses = append(nodesAddresses, node.Addr)
 			node_subscriptions_count = append(node_subscriptions_count, 0)
 		}
 
@@ -143,100 +159,123 @@ func main() {
 	total_channels := *channel_maximum - *channel_minimum + 1
 	subscriptions_per_node := total_channels / len(nodes)
 	total_subscriptions := total_channels * *subscribers_per_channel
-	fmt.Println(fmt.Sprintf("Total subcriptions: %d. Subscriptions per node %d", total_subscriptions, subscriptions_per_node))
+	total_messages := int64(total_subscriptions) * *messages_per_channel_subscriber
+	fmt.Println(fmt.Sprintf("Total subcriptions: %d. Subscriptions per node %d. Total messages: %d", total_subscriptions, subscriptions_per_node, total_messages))
+
+	// Launch routine to accumulate message count
+	//go accumulate()
 
 	for channel_id := *channel_minimum; channel_id <= *channel_maximum; channel_id++ {
 		for channel_subscriber_number := 1; channel_subscriber_number <= *subscribers_per_channel; channel_subscriber_number++ {
 			nodes_pos := channel_id % len(nodes)
-			node_subscriptions_count[nodes_pos] = node_subscriptions_count[nodes_pos] + 1
+			node_subscriptions_count[nodes_pos]++
 			addr := nodes[nodes_pos]
 			channel := fmt.Sprintf("%s%d", *subscribe_prefix, channel_id)
 			subscriberName := fmt.Sprintf("subscriber#%d-%s%d", channel_subscriber_number, *subscribe_prefix, channel_id)
 			wg.Add(1)
-			go subscriberRoutine(addr.Addr, subscriberName, channel, *printMessages, stopChan, &wg)
+			go subscriberRoutine(addr.Addr, subscriberName, channel, *messages_per_channel_subscriber, *printMessages, stopChan, &wg)
 		}
-
 	}
 
 	// listen for C-c
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
+	w := new(tabwriter.Writer)
 
 	tick := time.NewTicker(time.Duration(*client_update_tick) * time.Second)
 	var connections []radix.Conn
-	if updateCLI(nodes, connections, node_subscriptions_count, tick, c) {
+	closed, start_time, duration, totalMessages, messageRateTs := updateCLI(nodes, connections, node_subscriptions_count, tick, c, total_messages, w, *test_time)
+	messageRate := float64(totalMessages) / float64(duration.Seconds())
+	fmt.Fprint(w, fmt.Sprintf("#################################################\nTotal Duration %f Seconds\nMessage Rate %f\n#################################################\n", duration.Seconds(), messageRate))
+	fmt.Fprint(w, "\r\n")
+	w.Flush()
+
+	if strings.Compare(*json_out_file, "") != 0 {
+		res := testResult{
+			StartTime:             start_time.Unix(),
+			Duration:              duration.Seconds(),
+			MessageRate:           messageRate,
+			TotalMessages:         totalMessages,
+			TotalSubscriptions:    total_subscriptions,
+			ChannelMin:            *channel_minimum,
+			ChannelMax:            *channel_maximum,
+			SubscribersPerChannel: *subscribers_per_channel,
+			MessagesPerChannel:    *messages_per_channel_subscriber,
+			MessageRateTs:         messageRateTs,
+			OSSDistributedSlots:   *distributeSubscribers,
+			Addresses:             nodesAddresses,
+		}
+		file, err := json.MarshalIndent(res, "", " ")
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		err = ioutil.WriteFile(*json_out_file, file, 0644)
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+
+	if closed {
 		return
 	}
 
 	// tell the goroutine to stop
 	close(stopChan)
 	// and wait for them both to reply back
 	wg.Wait()
-
 }
 
-func updateCLI(nodes []radix.ClusterNode, connections []radix.Conn, node_subscriptions_count []int, tick *time.Ticker, c chan os.Signal) bool {
-	w := new(tabwriter.Writer)
+func updateCLI(nodes []radix.ClusterNode, connections []radix.Conn, node_subscriptions_count []int, tick *time.Ticker, c chan os.Signal, message_limit int64, w *tabwriter.Writer, test_time int) (bool, time.Time, time.Duration, uint64, []float64) {
+
+	start := time.Now()
+	prevTime := time.Now()
+	prevMessageCount := uint64(0)
+	messageRateTs := []float64{}
+
 	w.Init(os.Stdout, 25, 0, 1, ' ', tabwriter.AlignRight)
-	for idx, slot := range nodes {
-		c, err := radix.Dial("tcp", slot.Addr)
-		if err != nil {
-			panic(err)
-		}
-		fmt.Fprint(w, fmt.Sprintf("shard #%d\t", idx+1))
-		connections = append(connections, c)
-	}
+	fmt.Fprint(w, fmt.Sprintf("Test Time\tTotal Messages\t Message Rate \t"))
 	fmt.Fprint(w, "\n")
 	w.Flush()
 	for {
 		select {
 		case <-tick.C:
 			{
-				for idx, c := range connections {
-					var infoOutput string
-					e := c.Do(radix.FlatCmd(&infoOutput, "INFO", "CLIENTS"))
-					if e != nil {
-						fmt.Fprint(w, fmt.Sprintf("----\t"))
-					} else {
-						connected_clients_line := strings.TrimSuffix(strings.Split(infoOutput, "\r\n")[1], "\r\n")
-						i := strings.Index(connected_clients_line, ":")
-						shard_clients, eint := strconv.Atoi(connected_clients_line[i+1:])
-						if eint != nil {
-							fmt.Fprint(w, fmt.Sprintf("----\t"))
-						} else {
-							var pubsubList string
-							e := c.Do(radix.FlatCmd(&pubsubList, "CLIENT", "LIST", "TYPE","PUBSUB"))
-							if e != nil {
-								fmt.Fprint(w, fmt.Sprintf("----\t"))
-							}
-							subscribersList := strings.Split(pubsubList, "\n")
-							//fmt.Println(subscribersList)
-							expected_subscribers := node_subscriptions_count[idx]
-							number_subscribers := len(subscribersList)-1
-							others := shard_clients - 1 - number_subscribers
-
-							fmt.Fprint(w, fmt.Sprintf("%d (other %d sub %d==%d)\t", shard_clients, others, number_subscribers, expected_subscribers))
-						}
-
-					}
+				now := time.Now()
+				took := now.Sub(prevTime)
+				messageRate := float64(totalMessages-prevMessageCount) / float64(took.Seconds())
+				if prevMessageCount == 0 && totalMessages != 0 {
+					start = time.Now()
 				}
+				if totalMessages != 0 {
+					messageRateTs = append(messageRateTs, messageRate)
+				}
+				prevMessageCount = totalMessages
+				prevTime = now
+
+				fmt.Fprint(w, fmt.Sprintf("%.0f\t%d\t%.2f\t", time.Since(start).Seconds(), totalMessages, messageRate))
 				fmt.Fprint(w, "\r\n")
 				w.Flush()
+				if message_limit > 0 && totalMessages >= uint64(message_limit) {
+					return true, start, time.Since(start), totalMessages, messageRateTs
+				}
+				if test_time > 0 && time.Since(start) >= time.Duration(test_time*1000*1000*1000) && totalMessages != 0 {
+					return true, start, time.Since(start), totalMessages, messageRateTs
+				}
 
 				break
 			}
 
 		case <-c:
 			fmt.Println("received Ctrl-c - shutting down")
-			return true
+			return true, start, time.Since(start), totalMessages, messageRateTs
 		}
 	}
-	return false
+	return false, start, time.Since(start), totalMessages, messageRateTs
 }
 
 func checkClientOutputBufferLimitPubSub(nodes []radix.ClusterNode, client_output_buffer_limit_pubsub *string) {
 	for _, slot := range nodes {
-		//fmt.Println(slot)
 		conn, err := radix.Dial("tcp", slot.Addr)
 		if err != nil {
 			panic(err)