github.com/koko1123/flow-go-1@v0.29.6/engine/access/rpc/backend/connection_factory_test.go (about)

     1  package backend
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"strconv"
     8  	"strings"
     9  	"sync"
    10  	"testing"
    11  	"time"
    12  
    13  	lru "github.com/hashicorp/golang-lru"
    14  	"github.com/onflow/flow/protobuf/go/flow/access"
    15  	"github.com/onflow/flow/protobuf/go/flow/execution"
    16  	"github.com/stretchr/testify/assert"
    17  	testifymock "github.com/stretchr/testify/mock"
    18  	"google.golang.org/grpc"
    19  	"google.golang.org/grpc/codes"
    20  	"google.golang.org/grpc/status"
    21  
    22  	"github.com/koko1123/flow-go-1/engine/access/mock"
    23  	"github.com/koko1123/flow-go-1/module/metrics"
    24  	"github.com/koko1123/flow-go-1/utils/unittest"
    25  )
    26  
    27  func TestProxyAccessAPI(t *testing.T) {
    28  	// create a collection node
    29  	cn := new(collectionNode)
    30  	cn.start(t)
    31  	defer cn.stop(t)
    32  
    33  	req := &access.PingRequest{}
    34  	expected := &access.PingResponse{}
    35  	cn.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
    36  
    37  	// create the factory
    38  	connectionFactory := new(ConnectionFactoryImpl)
    39  	// set the collection grpc port
    40  	connectionFactory.CollectionGRPCPort = cn.port
    41  	// set metrics reporting
    42  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
    43  
    44  	proxyConnectionFactory := ProxyConnectionFactory{
    45  		ConnectionFactory: connectionFactory,
    46  		targetAddress:     cn.listener.Addr().String(),
    47  	}
    48  
    49  	// get a collection API client
    50  	client, conn, err := proxyConnectionFactory.GetAccessAPIClient("foo")
    51  	defer conn.Close()
    52  	assert.NoError(t, err)
    53  
    54  	ctx := context.Background()
    55  	// make the call to the collection node
    56  	resp, err := client.Ping(ctx, req)
    57  	assert.NoError(t, err)
    58  	assert.Equal(t, resp, expected)
    59  }
    60  
    61  func TestProxyExecutionAPI(t *testing.T) {
    62  	// create an execution node
    63  	en := new(executionNode)
    64  	en.start(t)
    65  	defer en.stop(t)
    66  
    67  	req := &execution.PingRequest{}
    68  	expected := &execution.PingResponse{}
    69  	en.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
    70  
    71  	// create the factory
    72  	connectionFactory := new(ConnectionFactoryImpl)
    73  	// set the execution grpc port
    74  	connectionFactory.ExecutionGRPCPort = en.port
    75  	// set metrics reporting
    76  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
    77  
    78  	proxyConnectionFactory := ProxyConnectionFactory{
    79  		ConnectionFactory: connectionFactory,
    80  		targetAddress:     en.listener.Addr().String(),
    81  	}
    82  
    83  	// get an execution API client
    84  	client, _, err := proxyConnectionFactory.GetExecutionAPIClient("foo")
    85  	assert.NoError(t, err)
    86  
    87  	ctx := context.Background()
    88  	// make the call to the execution node
    89  	resp, err := client.Ping(ctx, req)
    90  	assert.NoError(t, err)
    91  	assert.Equal(t, resp, expected)
    92  }
    93  
    94  func TestProxyAccessAPIConnectionReuse(t *testing.T) {
    95  	// create a collection node
    96  	cn := new(collectionNode)
    97  	cn.start(t)
    98  	defer cn.stop(t)
    99  
   100  	req := &access.PingRequest{}
   101  	expected := &access.PingResponse{}
   102  	cn.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
   103  
   104  	// create the factory
   105  	connectionFactory := new(ConnectionFactoryImpl)
   106  	// set the collection grpc port
   107  	connectionFactory.CollectionGRPCPort = cn.port
   108  	// set the connection pool cache size
   109  	cacheSize := 5
   110  	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
   111  		evictedValue.(*CachedClient).Close()
   112  	})
   113  	connectionFactory.ConnectionsCache = cache
   114  	connectionFactory.CacheSize = uint(cacheSize)
   115  	// set metrics reporting
   116  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
   117  
   118  	proxyConnectionFactory := ProxyConnectionFactory{
   119  		ConnectionFactory: connectionFactory,
   120  		targetAddress:     cn.listener.Addr().String(),
   121  	}
   122  
   123  	// get a collection API client
   124  	_, closer, err := proxyConnectionFactory.GetAccessAPIClient("foo")
   125  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 1)
   126  	assert.NoError(t, err)
   127  	assert.Nil(t, closer.Close())
   128  
   129  	var conn *grpc.ClientConn
   130  	res, ok := connectionFactory.ConnectionsCache.Get(proxyConnectionFactory.targetAddress)
   131  	assert.True(t, ok)
   132  	conn = res.(*CachedClient).ClientConn
   133  
   134  	// check if api client can be rebuilt with retrieved connection
   135  	accessAPIClient := access.NewAccessAPIClient(conn)
   136  	ctx := context.Background()
   137  	resp, err := accessAPIClient.Ping(ctx, req)
   138  	assert.NoError(t, err)
   139  	assert.Equal(t, resp, expected)
   140  }
   141  
   142  func TestProxyExecutionAPIConnectionReuse(t *testing.T) {
   143  	// create an execution node
   144  	en := new(executionNode)
   145  	en.start(t)
   146  	defer en.stop(t)
   147  
   148  	req := &execution.PingRequest{}
   149  	expected := &execution.PingResponse{}
   150  	en.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
   151  
   152  	// create the factory
   153  	connectionFactory := new(ConnectionFactoryImpl)
   154  	// set the execution grpc port
   155  	connectionFactory.ExecutionGRPCPort = en.port
   156  	// set the connection pool cache size
   157  	cacheSize := 5
   158  	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
   159  		evictedValue.(*CachedClient).Close()
   160  	})
   161  	connectionFactory.ConnectionsCache = cache
   162  	connectionFactory.CacheSize = uint(cacheSize)
   163  	// set metrics reporting
   164  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
   165  
   166  	proxyConnectionFactory := ProxyConnectionFactory{
   167  		ConnectionFactory: connectionFactory,
   168  		targetAddress:     en.listener.Addr().String(),
   169  	}
   170  
   171  	// get an execution API client
   172  	_, closer, err := proxyConnectionFactory.GetExecutionAPIClient("foo")
   173  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 1)
   174  	assert.NoError(t, err)
   175  	assert.Nil(t, closer.Close())
   176  
   177  	var conn *grpc.ClientConn
   178  	res, ok := connectionFactory.ConnectionsCache.Get(proxyConnectionFactory.targetAddress)
   179  	assert.True(t, ok)
   180  	conn = res.(*CachedClient).ClientConn
   181  
   182  	// check if api client can be rebuilt with retrieved connection
   183  	executionAPIClient := execution.NewExecutionAPIClient(conn)
   184  	ctx := context.Background()
   185  	resp, err := executionAPIClient.Ping(ctx, req)
   186  	assert.NoError(t, err)
   187  	assert.Equal(t, resp, expected)
   188  }
   189  
   190  // TestExecutionNodeClientTimeout tests that the execution API client times out after the timeout duration
   191  func TestExecutionNodeClientTimeout(t *testing.T) {
   192  
   193  	timeout := 10 * time.Millisecond
   194  
   195  	// create an execution node
   196  	en := new(executionNode)
   197  	en.start(t)
   198  	defer en.stop(t)
   199  
   200  	// setup the handler mock to not respond within the timeout
   201  	req := &execution.PingRequest{}
   202  	resp := &execution.PingResponse{}
   203  	en.handler.On("Ping", testifymock.Anything, req).After(timeout+time.Second).Return(resp, nil)
   204  
   205  	// create the factory
   206  	connectionFactory := new(ConnectionFactoryImpl)
   207  	// set the execution grpc port
   208  	connectionFactory.ExecutionGRPCPort = en.port
   209  	// set the execution grpc client timeout
   210  	connectionFactory.ExecutionNodeGRPCTimeout = timeout
   211  	// set the connection pool cache size
   212  	cacheSize := 5
   213  	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
   214  		evictedValue.(*CachedClient).Close()
   215  	})
   216  	connectionFactory.ConnectionsCache = cache
   217  	connectionFactory.CacheSize = uint(cacheSize)
   218  	// set metrics reporting
   219  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
   220  
   221  	// create the execution API client
   222  	client, _, err := connectionFactory.GetExecutionAPIClient(en.listener.Addr().String())
   223  	assert.NoError(t, err)
   224  
   225  	ctx := context.Background()
   226  	// make the call to the execution node
   227  	_, err = client.Ping(ctx, req)
   228  
   229  	// assert that the client timed out
   230  	assert.Equal(t, codes.DeadlineExceeded, status.Code(err))
   231  }
   232  
   233  // TestCollectionNodeClientTimeout tests that the collection API client times out after the timeout duration
   234  func TestCollectionNodeClientTimeout(t *testing.T) {
   235  
   236  	timeout := 10 * time.Millisecond
   237  
   238  	// create a collection node
   239  	cn := new(collectionNode)
   240  	cn.start(t)
   241  	defer cn.stop(t)
   242  
   243  	// setup the handler mock to not respond within the timeout
   244  	req := &access.PingRequest{}
   245  	resp := &access.PingResponse{}
   246  	cn.handler.On("Ping", testifymock.Anything, req).After(timeout+time.Second).Return(resp, nil)
   247  
   248  	// create the factory
   249  	connectionFactory := new(ConnectionFactoryImpl)
   250  	// set the collection grpc port
   251  	connectionFactory.CollectionGRPCPort = cn.port
   252  	// set the collection grpc client timeout
   253  	connectionFactory.CollectionNodeGRPCTimeout = timeout
   254  	// set the connection pool cache size
   255  	cacheSize := 5
   256  	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
   257  		evictedValue.(*CachedClient).Close()
   258  	})
   259  	connectionFactory.ConnectionsCache = cache
   260  	connectionFactory.CacheSize = uint(cacheSize)
   261  	// set metrics reporting
   262  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
   263  
   264  	// create the collection API client
   265  	client, _, err := connectionFactory.GetAccessAPIClient(cn.listener.Addr().String())
   266  	assert.NoError(t, err)
   267  
   268  	ctx := context.Background()
   269  	// make the call to the execution node
   270  	_, err = client.Ping(ctx, req)
   271  
   272  	// assert that the client timed out
   273  	assert.Equal(t, codes.DeadlineExceeded, status.Code(err))
   274  }
   275  
   276  // TestConnectionPoolFull tests that the LRU cache replaces connections when full
   277  func TestConnectionPoolFull(t *testing.T) {
   278  	// create a collection node
   279  	cn1, cn2, cn3 := new(collectionNode), new(collectionNode), new(collectionNode)
   280  	cn1.start(t)
   281  	cn2.start(t)
   282  	cn3.start(t)
   283  	defer cn1.stop(t)
   284  	defer cn2.stop(t)
   285  	defer cn3.stop(t)
   286  
   287  	req := &access.PingRequest{}
   288  	expected := &access.PingResponse{}
   289  	cn1.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
   290  	cn2.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
   291  	cn3.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
   292  
   293  	// create the factory
   294  	connectionFactory := new(ConnectionFactoryImpl)
   295  	// set the collection grpc port
   296  	connectionFactory.CollectionGRPCPort = cn1.port
   297  	// set the connection pool cache size
   298  	cacheSize := 2
   299  	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
   300  		evictedValue.(*CachedClient).Close()
   301  	})
   302  	connectionFactory.ConnectionsCache = cache
   303  	connectionFactory.CacheSize = uint(cacheSize)
   304  	// set metrics reporting
   305  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
   306  
   307  	cn1Address := "foo1:123"
   308  	cn2Address := "foo2:123"
   309  	cn3Address := "foo3:123"
   310  
   311  	// get a collection API client
   312  	_, _, err := connectionFactory.GetAccessAPIClient(cn1Address)
   313  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 1)
   314  	assert.NoError(t, err)
   315  
   316  	_, _, err = connectionFactory.GetAccessAPIClient(cn2Address)
   317  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 2)
   318  	assert.NoError(t, err)
   319  
   320  	_, _, err = connectionFactory.GetAccessAPIClient(cn1Address)
   321  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 2)
   322  	assert.NoError(t, err)
   323  
   324  	// Expecting to replace cn2 because cn1 was accessed more recently
   325  	_, _, err = connectionFactory.GetAccessAPIClient(cn3Address)
   326  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 2)
   327  	assert.NoError(t, err)
   328  
   329  	var hostnameOrIP string
   330  	hostnameOrIP, _, err = net.SplitHostPort(cn1Address)
   331  	assert.NoError(t, err)
   332  	grpcAddress1 := fmt.Sprintf("%s:%d", hostnameOrIP, connectionFactory.CollectionGRPCPort)
   333  	hostnameOrIP, _, err = net.SplitHostPort(cn2Address)
   334  	assert.NoError(t, err)
   335  	grpcAddress2 := fmt.Sprintf("%s:%d", hostnameOrIP, connectionFactory.CollectionGRPCPort)
   336  	hostnameOrIP, _, err = net.SplitHostPort(cn3Address)
   337  	assert.NoError(t, err)
   338  	grpcAddress3 := fmt.Sprintf("%s:%d", hostnameOrIP, connectionFactory.CollectionGRPCPort)
   339  
   340  	contains1 := connectionFactory.ConnectionsCache.Contains(grpcAddress1)
   341  	contains2 := connectionFactory.ConnectionsCache.Contains(grpcAddress2)
   342  	contains3 := connectionFactory.ConnectionsCache.Contains(grpcAddress3)
   343  
   344  	assert.True(t, contains1)
   345  	assert.False(t, contains2)
   346  	assert.True(t, contains3)
   347  }
   348  
   349  // TestConnectionPoolStale tests that a new connection will be established if the old one cached is stale
   350  func TestConnectionPoolStale(t *testing.T) {
   351  	// create a collection node
   352  	cn := new(collectionNode)
   353  	cn.start(t)
   354  	defer cn.stop(t)
   355  
   356  	req := &access.PingRequest{}
   357  	expected := &access.PingResponse{}
   358  	cn.handler.On("Ping", testifymock.Anything, req).Return(expected, nil)
   359  
   360  	// create the factory
   361  	connectionFactory := new(ConnectionFactoryImpl)
   362  	// set the collection grpc port
   363  	connectionFactory.CollectionGRPCPort = cn.port
   364  	// set the connection pool cache size
   365  	cacheSize := 5
   366  	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
   367  		evictedValue.(*CachedClient).Close()
   368  	})
   369  	connectionFactory.ConnectionsCache = cache
   370  	connectionFactory.CacheSize = uint(cacheSize)
   371  	// set metrics reporting
   372  	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
   373  
   374  	proxyConnectionFactory := ProxyConnectionFactory{
   375  		ConnectionFactory: connectionFactory,
   376  		targetAddress:     cn.listener.Addr().String(),
   377  	}
   378  
   379  	// get a collection API client
   380  	client, _, err := proxyConnectionFactory.GetAccessAPIClient("foo")
   381  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 1)
   382  	assert.NoError(t, err)
   383  	// close connection to simulate something "going wrong" with our stored connection
   384  	res, _ := connectionFactory.ConnectionsCache.Get(proxyConnectionFactory.targetAddress)
   385  
   386  	res.(*CachedClient).Close()
   387  
   388  	ctx := context.Background()
   389  	// make the call to the collection node (should fail, connection closed)
   390  	_, err = client.Ping(ctx, req)
   391  	assert.Error(t, err)
   392  
   393  	// re-access, should replace stale connection in cache with new one
   394  	_, _, _ = proxyConnectionFactory.GetAccessAPIClient("foo")
   395  	assert.Equal(t, connectionFactory.ConnectionsCache.Len(), 1)
   396  
   397  	var conn *grpc.ClientConn
   398  	res, ok := connectionFactory.ConnectionsCache.Get(proxyConnectionFactory.targetAddress)
   399  	assert.True(t, ok)
   400  	conn = res.(*CachedClient).ClientConn
   401  
   402  	// check if api client can be rebuilt with retrieved connection
   403  	accessAPIClient := access.NewAccessAPIClient(conn)
   404  	ctx = context.Background()
   405  	resp, err := accessAPIClient.Ping(ctx, req)
   406  	assert.NoError(t, err)
   407  	assert.Equal(t, resp, expected)
   408  }
   409  
   410  // node mocks a flow node that runs a GRPC server
   411  type node struct {
   412  	server   *grpc.Server
   413  	listener net.Listener
   414  	port     uint
   415  }
   416  
   417  func (n *node) setupNode(t *testing.T) {
   418  	n.server = grpc.NewServer()
   419  	listener, err := net.Listen("tcp4", unittest.DefaultAddress)
   420  	assert.NoError(t, err)
   421  	n.listener = listener
   422  	assert.Eventually(t, func() bool {
   423  		return !strings.HasSuffix(listener.Addr().String(), ":0")
   424  	}, time.Second*4, 10*time.Millisecond)
   425  
   426  	_, port, err := net.SplitHostPort(listener.Addr().String())
   427  	assert.NoError(t, err)
   428  	portAsUint, err := strconv.ParseUint(port, 10, 32)
   429  	assert.NoError(t, err)
   430  	n.port = uint(portAsUint)
   431  }
   432  
   433  func (n *node) start(t *testing.T) {
   434  	// using a wait group here to ensure the goroutine has started before returning. Otherwise,
   435  	// there's a race condition where the server is sometimes stopped before it has started
   436  	wg := sync.WaitGroup{}
   437  	wg.Add(1)
   438  	go func() {
   439  		wg.Done()
   440  		err := n.server.Serve(n.listener)
   441  		assert.NoError(t, err)
   442  	}()
   443  	unittest.RequireReturnsBefore(t, wg.Wait, 10*time.Millisecond, "could not start goroutine on time")
   444  }
   445  
   446  func (n *node) stop(t *testing.T) {
   447  	if n.server != nil {
   448  		n.server.Stop()
   449  	}
   450  }
   451  
   452  type executionNode struct {
   453  	node
   454  	handler *mock.ExecutionAPIServer
   455  }
   456  
   457  func (en *executionNode) start(t *testing.T) {
   458  	en.setupNode(t)
   459  	handler := new(mock.ExecutionAPIServer)
   460  	execution.RegisterExecutionAPIServer(en.server, handler)
   461  	en.handler = handler
   462  	en.node.start(t)
   463  }
   464  
   465  func (en *executionNode) stop(t *testing.T) {
   466  	en.node.stop(t)
   467  }
   468  
   469  type collectionNode struct {
   470  	node
   471  	handler *mock.AccessAPIServer
   472  }
   473  
   474  func (cn *collectionNode) start(t *testing.T) {
   475  	cn.setupNode(t)
   476  	handler := new(mock.AccessAPIServer)
   477  	access.RegisterAccessAPIServer(cn.server, handler)
   478  	cn.handler = handler
   479  	cn.node.start(t)
   480  }
   481  
   482  func (cn *collectionNode) stop(t *testing.T) {
   483  	cn.node.stop(t)
   484  }