[public-api] Implement connection pool with LRU cache

Author: easyCZ

Diff for: components/gitpod-protocol/go/gitpod-service.go

@@ -11,6 +11,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"net/url"
 	"sync"
@@ -23,6 +24,8 @@ import (
 
 // APIInterface wraps the
 type APIInterface interface {
+	io.Closer
+
 	GetOwnerToken(ctx context.Context, workspaceID string) (res string, err error)
 	AdminBlockUser(ctx context.Context, req *AdminBlockUserRequest) (err error)
 	GetLoggedInUser(ctx context.Context) (res *User, err error)

Diff for: components/gitpod-protocol/go/mock.go

@@ -33,7 +33,7 @@ require (
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/gorilla/websocket v1.5.0 // indirect
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 // indirect
-	github.com/hashicorp/golang-lru v0.5.1 // indirect
+	github.com/hashicorp/golang-lru v0.5.4 // indirect
 	github.com/heptiolabs/healthcheck v0.0.0-20211123025425-613501dd5deb // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect

Diff for: components/public-api-server/go.mod

@@ -7,10 +7,14 @@ package proxy
 import (
 	"context"
 	"fmt"
-	gitpod "github.com/gitpod-io/gitpod/gitpod-protocol"
-	"github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
 	"net/url"
 	"time"
+
+	"github.com/gitpod-io/gitpod/common-go/log"
+	gitpod "github.com/gitpod-io/gitpod/gitpod-protocol"
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
+
+	lru "github.com/hashicorp/golang-lru"
 )
 
 type ServerConnectionPool interface {
@@ -42,3 +46,69 @@ func (p *NoConnectionPool) Get(ctx context.Context, token string) (gitpod.APIInt
 
 	return server, nil
 }
+
+func NewConnectionPool(address *url.URL, poolSize int) (*ConnectionPool, error) {
+	cache, err := lru.NewWithEvict(poolSize, func(key, value interface{}) {
+		connectionPoolSize.Dec()
+
+		// We attempt to gracefully close the connection
+		conn, ok := value.(gitpod.APIInterface)
+		if !ok {
+			return
+		}
+
+		closeErr := conn.Close()
+		if closeErr != nil {
+			log.Log.WithError(closeErr).Warn("Failed to close connection to server.")
+		}
+	})
+	if err != nil {
+		return nil, fmt.Errorf("faield to create LRU cache: %w", err)
+	}
+
+	return &ConnectionPool{
+		cache: cache,
+		connConstructor: func(token string) (gitpod.APIInterface, error) {
+			return gitpod.ConnectToServer(address.String(), gitpod.ConnectToServerOpts{
+				// We want the connection to persist beyond the lifecycle of a single request
+				Context: context.Background(),
+				Token:   token,
+				Log:     log.Log,
+				CloseHandler: func(err error) {
+					cache.Remove(token)
+					connectionPoolSize.Dec()
+				},
+			})
+		},
+	}, nil
+
+}
+
+type ConnectionPool struct {
+	connConstructor func(token string) (gitpod.APIInterface, error)
+
+	// cache stores token to connection mapping
+	cache *lru.Cache
+}
+
+func (p *ConnectionPool) Get(ctx context.Context, token string) (gitpod.APIInterface, error) {
+	cached, found := p.cache.Get(token)
+	reportCacheOutcome(found)
+	if found {
+		connectionPoolCacheOutcome.WithLabelValues("true").Inc()
+		conn, ok := cached.(*gitpod.APIoverJSONRPC)
+		if ok {
+			return conn, nil
+		}
+	}
+
+	conn, err := p.connConstructor(token)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create new connection to server: %w", err)
+	}
+
+	p.cache.Add(token, conn)
+	connectionPoolSize.Inc()
+
+	return conn, nil
+}

Diff for: components/public-api-server/go.sum

@@ -0,0 +1,44 @@
+// Copyright (c) 2022 Gitpod GmbH. All rights reserved.
+// Licensed under the GNU Affero General Public License (AGPL).
+// See License-AGPL.txt in the project root for license information.
+
+package proxy
+
+import (
+	"context"
+	"testing"
+
+	gitpod "github.com/gitpod-io/gitpod/gitpod-protocol"
+	"github.com/golang/mock/gomock"
+	lru "github.com/hashicorp/golang-lru"
+	"github.com/stretchr/testify/require"
+)
+
+func TestConnectionPool(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	srv := gitpod.NewMockAPIInterface(ctrl)
+
+	cache, err := lru.New(2)
+	require.NoError(t, err)
+	pool := &ConnectionPool{
+		cache: cache,
+		connConstructor: func(token string) (gitpod.APIInterface, error) {
+			return srv, nil
+		},
+	}
+
+	_, err = pool.Get(context.Background(), "foo")
+	require.NoError(t, err)
+	require.Equal(t, 1, pool.cache.Len())
+
+	_, err = pool.Get(context.Background(), "bar")
+	require.NoError(t, err)
+	require.Equal(t, 2, pool.cache.Len())
+
+	_, err = pool.Get(context.Background(), "baz")
+	require.NoError(t, err)
+	require.Equal(t, 2, pool.cache.Len(), "must keep only last two connectons")
+	require.True(t, pool.cache.Contains("bar"))
+	require.True(t, pool.cache.Contains("baz"))
+}

Diff for: components/public-api-server/pkg/proxy/conn.go

@@ -5,20 +5,45 @@
 package proxy
 
 import (
-	"github.com/prometheus/client_golang/prometheus"
+	"strconv"
 	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 func reportConnectionDuration(d time.Duration) {
 	proxyConnectionCreateDurationSeconds.Observe(d.Seconds())
 }
 
-var proxyConnectionCreateDurationSeconds = prometheus.NewHistogram(prometheus.HistogramOpts{
-	Namespace: "gitpod",
-	Name:      "public_api_proxy_connection_create_duration_seconds",
-	Help:      "Histogram of connection time in seconds",
-})
+var (
+	proxyConnectionCreateDurationSeconds = prometheus.NewHistogram(prometheus.HistogramOpts{
+		Namespace: "gitpod",
+		Subsystem: "public_api",
+		Name:      "proxy_connection_create_duration_seconds",
+		Help:      "Histogram of connection time in seconds",
+	})
+
+	connectionPoolSize = prometheus.NewGauge(prometheus.GaugeOpts{
+		Namespace: "gitpod",
+		Subsystem: "public_api",
+		Name:      "proxy_connection_pool_size",
+		Help:      "Gauge of connections in connection pool",
+	})
+
+	connectionPoolCacheOutcome = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Namespace: "gitpod",
+		Subsystem: "public_api",
+		Name:      "proxy_connection_pool_cache_outcomes_total",
+		Help:      "Counter of cachce accesses",
+	}, []string{"hit"})
+)
 
 func RegisterMetrics(registry *prometheus.Registry) {
 	registry.MustRegister(proxyConnectionCreateDurationSeconds)
+	registry.MustRegister(connectionPoolSize)
+	registry.MustRegister(connectionPoolCacheOutcome)
+}
+
+func reportCacheOutcome(hit bool) {
+	connectionPoolCacheOutcome.WithLabelValues(strconv.FormatBool(hit)).Inc()
 }

Diff for: components/public-api-server/pkg/proxy/conn_test.go

@@ -35,7 +35,10 @@ func Start(logger *logrus.Entry, version string, cfg *config.Configuration) erro
 		return fmt.Errorf("failed to parse Gitpod API URL: %w", err)
 	}
 
-	connPool := &proxy.NoConnectionPool{ServerAPI: gitpodAPI}
+	connPool, err := proxy.NewConnectionPool(gitpodAPI, 3000)
+	if err != nil {
+		return fmt.Errorf("failed to setup connection pool: %w", err)
+	}
 
 	srv, err := baseserver.New("public_api_server",
 		baseserver.WithLogger(logger),
@@ -82,7 +85,6 @@ func register(srv *baseserver.Server, connPool proxy.ServerConnectionPool) error
 	proxy.RegisterMetrics(srv.MetricsRegistry())
 
 	connectMetrics := NewConnectMetrics()
-
 	err := connectMetrics.Register(srv.MetricsRegistry())
 	if err != nil {
 		return err