go.uber.org/yarpc@v1.72.1/encoding/thrift/benchmark_test.go (about)

     1  package thrift_test
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"net"
     7  	"strconv"
     8  	"strings"
     9  	"sync"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  	"go.uber.org/yarpc"
    16  	"go.uber.org/yarpc/api/transport"
    17  	"go.uber.org/yarpc/internal/examples/thrift-keyvalue/keyvalue/kv"
    18  	"go.uber.org/yarpc/internal/examples/thrift-keyvalue/keyvalue/kv/keyvalueclient"
    19  	"go.uber.org/yarpc/internal/examples/thrift-keyvalue/keyvalue/kv/keyvalueserver"
    20  	"go.uber.org/yarpc/transport/tchannel"
    21  )
    22  
    23  const (
    24  	_kvServer = "callee"
    25  	_kvClient = "caller"
    26  )
    27  
    28  func BenchmarkThriftClientCallNormalDist(b *testing.B) {
    29  	handler := &keyValueHandler{}
    30  	serverAddr := newKeyValServer(b, handler)
    31  
    32  	clientNoReuse := newKeyValueClient(b, serverAddr, false)
    33  	clientWithReuse := newKeyValueClient(b, serverAddr, true)
    34  
    35  	// Create a normal distribution
    36  	// deviation 10k, mean 3KB, minimum 0, maximum 2MB
    37  	g := createNormalDistribution(3*1024, 10_000, 0, 2*1024*1024)
    38  
    39  	var samples []string
    40  	for i := 0; i < 10000; i++ {
    41  		key := "foo" + strconv.FormatInt(int64(i), 10)
    42  		length := g()
    43  		value := generateRandomString(length)
    44  		samples = append(samples, value)
    45  		handler.SetValue(context.Background(), &key, &value)
    46  	}
    47  
    48  	b.ResetTimer()
    49  
    50  	b.Run("with_buffer_pool", func(b *testing.B) {
    51  		b.ReportAllocs()
    52  		for i := 0; i < b.N; i++ {
    53  			offset := i % len(samples)
    54  			key := "foo" + strconv.FormatInt(int64(offset), 10)
    55  			value := samples[i%len(samples)]
    56  			callGetter(b, clientWithReuse, key, value)
    57  		}
    58  	})
    59  
    60  	b.Run("without_buffer_pool", func(b *testing.B) {
    61  		b.ReportAllocs()
    62  		for i := 0; i < b.N; i++ {
    63  			offset := i % len(samples)
    64  			key := "foo" + strconv.FormatInt(int64(offset), 10)
    65  			value := samples[i%len(samples)]
    66  			callGetter(b, clientNoReuse, key, value)
    67  		}
    68  	})
    69  }
    70  
    71  func generateRandomString(len int) string {
    72  	var sb strings.Builder
    73  	for i := 0; i < len; i++ {
    74  		c := 'a' + rand.Intn('z'-'a')
    75  		sb.WriteByte(byte(c))
    76  	}
    77  	return sb.String()
    78  }
    79  
    80  func callGetter(b *testing.B, client keyvalueclient.Interface, key string, want string) {
    81  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
    82  	defer cancel()
    83  
    84  	got, err := client.GetValue(ctx, &key)
    85  	require.NoError(b, err)
    86  	require.Equal(b, want, got)
    87  }
    88  
    89  type keyValueHandler struct {
    90  	items sync.Map
    91  }
    92  
    93  func (h *keyValueHandler) GetValue(ctx context.Context, key *string) (string, error) {
    94  	if v, ok := h.items.Load(*key); ok {
    95  		return v.(string), nil
    96  	}
    97  	return "", &kv.ResourceDoesNotExist{Key: *key}
    98  }
    99  
   100  func (h *keyValueHandler) SetValue(ctx context.Context, key *string, value *string) error {
   101  	h.items.Store(*key, *value)
   102  	return nil
   103  }
   104  
   105  func newKeyValServer(t testing.TB, handler keyvalueserver.Interface) string {
   106  	listen, err := net.Listen("tcp", "127.0.0.1:0")
   107  	require.NoError(t, err)
   108  
   109  	trans, err := tchannel.NewTransport(
   110  		tchannel.ServiceName(_kvServer),
   111  		tchannel.Listener(listen))
   112  	require.NoError(t, err)
   113  
   114  	inbound := trans.NewInbound()
   115  	addr := listen.Addr().String()
   116  
   117  	dispatcher := yarpc.NewDispatcher(yarpc.Config{
   118  		Name:     _kvServer,
   119  		Inbounds: yarpc.Inbounds{inbound},
   120  	})
   121  
   122  	dispatcher.Register(keyvalueserver.New(handler))
   123  	require.NoError(t, dispatcher.Start(), "could not start server dispatcher")
   124  
   125  	t.Cleanup(func() { assert.NoError(t, dispatcher.Stop(), "could not stop dispatcher") })
   126  
   127  	return addr
   128  }
   129  
   130  func newKeyValueClient(t testing.TB, serverAddr string, enableBufferReuse bool) keyvalueclient.Interface {
   131  	trans, err := tchannel.NewTransport(tchannel.ServiceName(_kvClient))
   132  	require.NoError(t, err)
   133  	out := trans.NewSingleOutbound(serverAddr, tchannel.WithReuseBuffer(enableBufferReuse))
   134  
   135  	dispatcher := yarpc.NewDispatcher(yarpc.Config{
   136  		Name: _kvClient,
   137  		Outbounds: map[string]transport.Outbounds{
   138  			_kvServer: {
   139  				ServiceName: _kvServer,
   140  				Unary:       out,
   141  			},
   142  		},
   143  	})
   144  
   145  	client := keyvalueclient.New(dispatcher.ClientConfig(_kvServer))
   146  	require.NoError(t, dispatcher.Start(), "could not start client dispatcher")
   147  
   148  	t.Cleanup(func() { assert.NoError(t, dispatcher.Stop(), "could not stop dispatcher") })
   149  	return client
   150  }
   151  
   152  func createNormalDistribution(mean, deviation float64, minimum, maximum int) func() int {
   153  	r := rand.New(rand.NewSource(time.Now().UnixNano()))
   154  	return func() int {
   155  		n := int(r.NormFloat64()*mean + deviation)
   156  
   157  		// 0 <= n <= 2MB
   158  		n = max(n, minimum)
   159  		n = min(n, maximum)
   160  
   161  		return n
   162  	}
   163  }