Experimental download-style streaming responses

Makefile

@@ -11,11 +11,11 @@ generate:
 	PATH=$(CURDIR)/_tools/bin:$(PATH) GOBIN="${PWD}/bin" go install -v ./protoc-gen-...
 	$(RETOOL) do go generate ./...
 
-test_all: setup test_core test_clients
+test_all: setup test_core test_clients test_example
 
 test_core: generate
 # $(RETOOL) do errcheck -blank ./internal/twirptest
-	go test -race $(shell go list ./... | grep -v /vendor/ | grep -v /_tools/)
+	go test -race $(shell go list ./... | grep -v /vendor/ | grep -v /_tools/ | grep -v /example/)
 
 test_clients: test_go_client test_python_client
 
@@ -25,6 +25,9 @@ test_go_client: generate build/clientcompat build/gocompat
 test_python_client: generate build/clientcompat build/pycompat
 	./build/clientcompat -client ./build/pycompat
 
+test_example: generate
+	go test -race -bench=. $(shell go list ./example/...)
+
 setup:
 	./install_proto.bash
 	GOPATH=$(CURDIR)/_tools go install github.com/twitchtv/retool/...

example/cmd/client/main.go

@@ -15,21 +15,25 @@ package main
 
 import (
 	"context"
-	"fmt"
 	"log"
 	"net/http"
+	"time"
 
 	"github.com/twitchtv/twirp"
 	"github.com/twitchtv/twirp/example"
 )
 
 func main() {
-	client := example.NewHaberdasherJSONClient("http://localhost:8080", &http.Client{})
+	client := example.NewHaberdasherProtobufClient("http://localhost:8080", &http.Client{})
 
 	var (
 		hat *example.Hat
 		err error
 	)
+
+	//
+	// Call the MakeHat rpc
+	//
 	for i := 0; i < 5; i++ {
 		hat, err = client.MakeHat(context.Background(), &example.Size{Inches: 12})
 		if err != nil {
@@ -43,6 +47,43 @@ func main() {
 			// This was some fatal error!
 			log.Fatal(err)
 		}
+		break
+	}
+	log.Println(`Response from MakeHat:`)
+	log.Printf("\t%+v\n", hat)
+
+	//
+	// Call the MakeHats streaming rpc
+	//
+	const (
+		printEvery = 50000
+		quantity   = int32(300000)
+	)
+	reqSentAt := time.Now()
+	hatStream, err := client.MakeHats(
+		context.Background(),
+		&example.MakeHatsReq{Inches: 12, Quantity: quantity},
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+	log.Printf("Response from MakeHats:\n")
+	ii := 1
+	printResults := func() {
+		took := time.Now().Sub(reqSentAt)
+		khps := float64(ii-1) / took.Seconds() / 1000
+		log.Printf("Received %.1f kHats per second (%d hats in %f seconds)\n", khps, ii-1, took.Seconds())
+	}
+	for hatOrErr := range hatStream {
+		if hatOrErr.Err != nil {
+			printResults()
+			log.Fatal(hatOrErr.Err)
+		}
+		if ii%printEvery == 0 {
+			khps := float64(ii) / time.Now().Sub(reqSentAt).Seconds() / 1000
+			log.Printf("\t[%4.1f khps] %6d %+v\n", khps, ii, hatOrErr.Msg)
+		}
+		ii++
 	}
-	fmt.Printf("%+v", hat)
+	printResults()
 }

example/cmd/server/main.go

@@ -27,17 +27,50 @@ import (
 
 type randomHaberdasher struct{}
 
-func (h *randomHaberdasher) MakeHat(ctx context.Context, size *example.Size) (*example.Hat, error) {
-	if size.Inches <= 0 {
-		return nil, twirp.InvalidArgumentError("Inches", "I can't make a hat that small!")
+var (
+	errTooSmall         = twirp.InvalidArgumentError("Inches", "I can't make hats that small!")
+	errNegativeQuantity = twirp.InvalidArgumentError("Quantity", "I can't make a negative quantity of hats!")
+)
+
+func newRandomHat(inches int32) (*example.Hat, error) {
+	if inches <= 0 {
+		return nil, errTooSmall
 	}
 	return &example.Hat{
-		Size:  size.Inches,
-		Color: []string{"white", "black", "brown", "red", "blue"}[rand.Intn(4)],
-		Name:  []string{"bowler", "baseball cap", "top hat", "derby"}[rand.Intn(3)],
+		Size:  inches,
+		Color: []string{"white", "black", "brown", "red", "blue"}[rand.Intn(5)],
+		Name:  []string{"bowler", "baseball cap", "top hat", "derby"}[rand.Intn(4)],
 	}, nil
 }
 
+func (h *randomHaberdasher) MakeHat(ctx context.Context, size *example.Size) (*example.Hat, error) {
+	return newRandomHat(size.Inches)
+}
+
+func (h *randomHaberdasher) MakeHats(ctx context.Context, req *example.MakeHatsReq) (<-chan example.HatOrError, error) {
+	if req.Quantity < 0 {
+		return nil, errNegativeQuantity
+	}
+	// Normally we'd validate Inches here as well, but we let it fall through to error on newRandomHat to demonstrate mid-stream errors
+	// if req.Inches <= 0 {
+	// 	return nil, errTooSmall
+	// }
+
+	ch := make(chan example.HatOrError, 100) // NB: the size of this buffer can make a big difference!
+	go func() {
+		defer close(ch)
+		for ii := int32(0); ii < req.Quantity; ii++ {
+			hat, err := newRandomHat(req.Inches)
+			select {
+			case <-ctx.Done():
+				return
+			case ch <- example.HatOrError{Msg: hat, Err: err}:
+			}
+		}
+	}()
+	return ch, nil
+}
+
 func main() {
 	hook := statsd.NewStatsdServerHooks(LoggingStatter{os.Stderr})
 	server := example.NewHaberdasherServer(&randomHaberdasher{}, hook)