Introduce support for pipelined requests, key prefix stripping, custom IPs, and Unix sockets

Author: adesjardins-stripe

cmd/basic/main.go

@@ -3,12 +3,13 @@ package main
 import (
 	"flag"
 	"fmt"
-	pcgr "github.com/dgryski/go-pcgr"
-	mct "github.com/dormando/mctester"
 	"math/rand"
 	"os"
 	"runtime/pprof"
 	"time"
+
+	"github.com/dgryski/go-pcgr"
+	mct "github.com/memcached/mctester"
 )
 
 var cpuprofile = flag.String("cpuprofile", "", "dump cpu profile to file")
@@ -31,6 +32,10 @@ func main() {
 	zipfV := flag.Float64("zipfV", float64(*keySpace/2), "zipf V value (pull below this number")
 	valueSize := flag.Uint("valuesize", 1000, "size of value (in bytes) to store on miss")
 	clientFlags := flag.Uint("clientflags", 0, "(32bit unsigned) client flag bits to set on miss")
+	pipelines := flag.Uint("pipelines", 1, "(32bit unsigned) stack this many GET requests into the same syscall.")
+	server := flag.String("server", "127.0.0.1:11211", "ip and port to connect to")
+	socket := flag.String("socket", "", "domain socket to connect to")
+	stripKeyPrefix := flag.Bool("stripkeyprefix", false, "strip key prefix before comparing with response.")
 
 	flag.Parse()
 
@@ -53,7 +58,10 @@ func main() {
 	*/
 
 	bl := &BasicLoader{
-		servers:               []string{"127.0.0.1:11211"},
+		servers:               []string{*server},
+		socket:                *socket,
+		pipelines:             *pipelines,
+		stripKeyPrefix:        *stripKeyPrefix,
 		desiredConnCount:      *connCount,
 		requestsPerSleep:      *reqPerSleep,
 		requestBundlesPerConn: *reqBundlePerConn,
@@ -100,6 +108,9 @@ func main() {
 // - variances: how often to change item sizes
 type BasicLoader struct {
 	servers               []string
+	socket                string
+	pipelines             uint
+	stripKeyPrefix        bool
 	stopAfter             time.Time
 	desiredConnCount      int
 	requestsPerSleep      int
@@ -141,8 +152,8 @@ func (l *BasicLoader) Run() {
 
 func (l *BasicLoader) Timer(tag string, start time.Time) {
 	duration := time.Since(start)
-	if duration > time.Millisecond * 10 {
-		fmt.Printf("%s [%d]\n", tag, int64(time.Since(start) / time.Microsecond))
+	if duration > time.Millisecond*10 {
+		fmt.Printf("%s [%d]\n", tag, int64(time.Since(start)/time.Microsecond))
 	}
 }
 
@@ -152,7 +163,7 @@ func (l *BasicLoader) Timer(tag string, start time.Time) {
 func (l *BasicLoader) Worker(doneChan chan<- int) {
 	// FIXME: selector.
 	host := l.servers[0]
-	mc := mct.NewClient(host)
+	mc := mct.NewClient(host, l.socket, l.pipelines, l.keyPrefix, l.stripKeyPrefix)
 	bundles := l.requestBundlesPerConn
 
 	rs := pcgr.New(time.Now().UnixNano(), 0)
@@ -185,15 +196,13 @@ func (l *BasicLoader) Worker(doneChan chan<- int) {
 			// Could also re-seed it twice, pull once Intn for length,
 			// re-seed, then again for key space.
 
-			keyLen := l.keyLength
 			if l.useZipf {
 				subRS.Seed(int64(zipRS.Uint64()))
 			} else {
 				subRS.Seed(int64(randR.Intn(l.keySpace)))
 			}
-			// TODO: might be nice to pass (by ref?) prefix in here to make
-			// use of string.Builder.
-			key := l.keyPrefix + mct.RandString(&subRS, keyLen)
+
+			key := mct.RandString(&subRS, l.keyLength, l.keyPrefix)
 			// chance we issue a delete instead.
 			delChance := randR.Intn(1000)
 			if l.deletePercent != 0 && delChance < l.deletePercent {

cmd/server/loader_basic.go

@@ -2,35 +2,42 @@ package main
 
 import (
 	"fmt"
-	pcgr "github.com/dgryski/go-pcgr"
-	mct "github.com/dormando/mctester"
 	"math/rand"
 	"time"
+
+	"github.com/dgryski/go-pcgr"
+	mct "github.com/memcached/mctester"
 )
 
 // Basic persistent load test, using text protocol:
 type BasicLoader struct {
-	Servers               []string `json:"servers"`
-	stopAfter             time.Time
-	DesiredConnCount      int `json:"conncount"`
-	RequestsPerSleep      int `json:"reqpersleep"`
-	RequestBundlesPerConn int `json:"reqbundlesperconn"`
+	Servers               []string      `json:"servers"`
+	Socket                string        `json:"socket"`
+	Pipelines             uint          `json:"pipelines"`
+	StripKeyPrefix        bool          `json:"stripkeyprefix"`
+	DesiredConnCount      int           `json:"conncount"`
+	RequestsPerSleep      int           `json:"reqpersleep"`
+	RequestBundlesPerConn int           `json:"reqbundlesperconn"`
 	SleepPerBundle        time.Duration `json:"sleepperbundle"`
-	DeletePercent         int `json:"deletepercent"`
-	KeyLength             int `json:"keylength"`
-	KeyPrefix             string `json:"keyprefix"`
-	KeySpace              int `json:"keyspace"`
-	KeyTTL                uint `json:"keyttl"`
-	UseZipf               bool `json:"zipf"`
-	ZipfS                 float64 `json:"zipfS"` // (> 1, generally 1.01-2) pulls the power curve toward 0)
-	ZipfV                 float64 `json:"zipfV"` // v (< KeySpace) puts the main part of the curve before this number
-	ValueSize             uint `json:"valuesize"`
-	ClientFlags           uint `json:"clientflags"`
+	DeletePercent         int           `json:"deletepercent"`
+	KeyLength             int           `json:"keylength"`
+	KeyPrefix             string        `json:"keyprefix"`
+	KeySpace              int           `json:"keyspace"`
+	KeyTTL                uint          `json:"keyttl"`
+	UseZipf               bool          `json:"zipf"`
+	ZipfS                 float64       `json:"zipfS"` // (> 1, generally 1.01-2) pulls the power curve toward 0)
+	ZipfV                 float64       `json:"zipfV"` // v (< KeySpace) puts the main part of the curve before this number
+	ValueSize             uint          `json:"valuesize"`
+	ClientFlags           uint          `json:"clientflags"`
+	stopAfter             time.Time
 }
 
 func newBasicLoader() *BasicLoader {
 	return &BasicLoader{
 		Servers:               []string{"127.0.0.1:11211"},
+		Socket:                "",
+		Pipelines:             1,
+		StripKeyPrefix:        false,
 		DesiredConnCount:      1,
 		RequestsPerSleep:      1,
 		RequestBundlesPerConn: 1,
@@ -116,7 +123,7 @@ func runBasicLoader(Update <-chan interface{}, worker interface{}) {
 func basicWorker(id int, doneChan chan<- int, updateChan <-chan *BasicLoader, l *BasicLoader) {
 	// TODO: server selector.
 	host := l.Servers[0]
-	mc := mct.NewClient(host)
+	mc := mct.NewClient(host, l.Socket, l.Pipelines, l.KeyPrefix, l.StripKeyPrefix)
 	bundles := l.RequestBundlesPerConn
 
 	rs := pcgr.New(time.Now().UnixNano(), 0)
@@ -149,15 +156,13 @@ func basicWorker(id int, doneChan chan<- int, updateChan <-chan *BasicLoader, l
 			// Could also re-seed it twice, pull once Intn for length,
 			// re-seed, then again for key space.
 
-			keyLen := l.KeyLength
 			if l.UseZipf {
 				subRS.Seed(int64(zipRS.Uint64()))
 			} else {
 				subRS.Seed(int64(randR.Intn(l.KeySpace)))
 			}
-			// TODO: might be nice to pass (by ref?) prefix in here to make
-			// use of string.Builder.
-			key := l.KeyPrefix + mct.RandString(&subRS, keyLen)
+
+			key := mct.RandString(&subRS, l.KeyLength, l.KeyPrefix)
 			// chance we issue a delete instead.
 			if l.DeletePercent != 0 && randR.Intn(1000) < l.DeletePercent {
 				_, err := mc.Delete(key)

protocol.go

@@ -17,6 +17,7 @@ import (
 	"io"
 	"net"
 	"strconv"
+	"strings"
 	"time"
 )
 
@@ -36,7 +37,13 @@ type mcConn struct {
 }
 
 func (c *Client) connectToMc() (*mcConn, error) {
-	conn, err := net.DialTimeout("tcp", c.Host, c.ConnectTimeout)
+	var conn net.Conn
+	var err error
+	if c.socket != "" {
+		conn, err = net.DialTimeout("unix", c.socket, c.ConnectTimeout)
+	} else {
+		conn, err = net.DialTimeout("tcp", c.Host, c.ConnectTimeout)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -50,19 +57,27 @@ type Client struct {
 	// read or write timeout
 	NetTimeout time.Duration
 	Host       string
+	socket     string
 	cn         *mcConn
 	WBufSize   int
 	RBufSize   int
 	// any necessary locks? channels?
 	// binprot structure cache.
-	binpkt *packet
-	opaque uint32 // just for binprot?
+	binpkt         *packet
+	opaque         uint32 // just for binprot?
+	pipelines      int
+	keyPrefix      string
+	stripKeyPrefix bool
 }
 
-func NewClient(host string) (client *Client) {
+func NewClient(host string, socket string, pipelines uint, keyPrefix string, stripKeyPrefix bool) (client *Client) {
 	client = &Client{
-		Host:   host,
-		binpkt: &packet{},
+		Host:           host,
+		socket:         socket,
+		pipelines:      int(pipelines),
+		keyPrefix:      keyPrefix,
+		stripKeyPrefix: stripKeyPrefix,
+		binpkt:         &packet{},
 	}
 	//client.rs = rand.NewSource(time.Now().UnixNano())
 	return client
@@ -319,56 +334,72 @@ func (c *Client) MetaDebug(key string) (err error) {
 //////////////////////////////////////////////
 
 func (c *Client) Get(key string) (flags uint64, value []byte, code McCode, err error) {
+	pipelines := c.pipelines
+	// Expected key from response
+	respKey := key
+	if c.stripKeyPrefix {
+		respKey = strings.TrimPrefix(key, c.keyPrefix)
+	}
+
 	err = c.runNow(key, len(key)+6, func() error {
 		b := c.cn.b
-		b.WriteString("get ")
-		b.WriteString(key)
-		b.WriteString("\r\n")
-		err = b.Flush()
-
-		if err != nil {
-			return err
+		for i := 0; i < pipelines; i++ {
+			b.WriteString("get ")
+			b.WriteString(key)
+			b.WriteString("\r\n")
 		}
-
-		line, err := b.ReadBytes('\n')
+		err = b.Flush()
 		if err != nil {
 			return err
 		}
 
-		if bytes.Equal(line, []byte("END\r\n")) {
-			code = McMISS
-		} else {
-			parts := bytes.Split(line[:len(line)-2], []byte(" "))
-			if !bytes.Equal(parts[0], []byte("VALUE")) {
-				// TODO: This should look for ERROR/SERVER_ERROR/etc
-				return ErrUnexpectedResponse
-			}
-			if len(parts) != 4 {
-				return ErrUnexpectedResponse
-			}
-			if !bytes.Equal(parts[1], []byte(key)) {
-				// FIXME: how do we embed the received vs expected in here?
-				// use the brand-new golang error wrapping thing?
-				return ErrKeyDoesNotMatch
-			}
-			flags, _ = ParseUint(parts[2])
-			size, _ := ParseUint(parts[3])
-
-			value = make([]byte, size+2)
-			_, err := io.ReadFull(b, value)
+		for i := 0; i < pipelines; i++ {
+			line, err := b.ReadBytes('\n')
 			if err != nil {
 				return err
 			}
 
-			if !bytes.Equal(value[len(value)-2:], []byte("\r\n")) {
-				return ErrCorruptValue
-			}
-			code = McHIT
-			value = value[:size]
-
-			line, err = b.ReadBytes('\n')
-			if !bytes.Equal(line, []byte("END\r\n")) {
-				return ErrUnexpectedResponse
+			if bytes.Equal(line, []byte("END\r\n")) {
+				code = McMISS
+			} else {
+				parts := bytes.Split(line[:len(line)-2], []byte(" "))
+				if !bytes.Equal(parts[0], []byte("VALUE")) {
+					// TODO: This should look for ERROR/SERVER_ERROR/etc
+					fmt.Print("Unexpected Response: ", string(line), "\n")
+					continue
+				}
+				if len(parts) != 4 {
+					fmt.Print("Unexpected Response: ", "parts not 4", "\n")
+					continue
+				}
+				if !bytes.Equal(parts[1], []byte(respKey)) {
+					fmt.Print("Unmatched Key: ", string(parts[1]), " and ", respKey, "\n")
+					// FIXME: how do we embed the received vs expected in here?
+					// use the brand-new golang error wrapping thing?
+					continue
+				}
+				flags, _ = ParseUint(parts[2])
+				size, _ := ParseUint(parts[3])
+
+				value = make([]byte, size+2)
+				_, err := io.ReadFull(b, value)
+				if err != nil {
+					fmt.Print("io ReadFull error, return", "\n")
+					return err
+				}
+
+				if !bytes.Equal(value[len(value)-2:], []byte("\r\n")) {
+					fmt.Print("Unmatched Value", "\n")
+					continue
+				}
+				code = McHIT
+				value = value[:size]
+
+				line, err = b.ReadBytes('\n')
+				if !bytes.Equal(line, []byte("END\r\n")) {
+					fmt.Print("Unmatched Reponse: ", string(line), " is not END\r\n")
+					continue
+				}
 			}
 		}
 

protocol_test.go

@@ -11,9 +11,13 @@ import (
 
 // tests expect a recent memcached to be running at this address.
 const hostname = "127.0.0.1:11211"
+const socket = ""
+const pipelines = 1
+const keyPrefix = "mctester:"
+const stripKeyPrefix = false
 
 func newcli() *Client {
-	mc := NewClient(hostname)
+	mc := NewClient(hostname, socket, pipelines, keyPrefix, stripKeyPrefix)
 	mc.ConnectTimeout = 3 * time.Second
 	mc.NetTimeout = time.Second
 	mc.WBufSize = 64 * 1024

support.go

@@ -25,9 +25,10 @@ const (
 
 // randomized keys!
 // TODO: is sb reusable?
-func RandString(src rand.Source, n int) string {
+func RandString(src rand.Source, n int, prefix string) string {
 	sb := strings.Builder{}
-	sb.Grow(n)
+	sb.Grow(len(prefix) + n)
+	sb.WriteString(prefix)
 	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
 	for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
 		if remain == 0 {