-
Notifications
You must be signed in to change notification settings - Fork 3.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
ccl/sqlproxyccl: add basic Balancer component for load balancing
This commit adds a basic Balancer component that is responsible for load balancing connections across SQL pods. This commit updates the connector component to invoke methods on the balancer when selecting a pod instead of the tenant directory. For now, the balancer does nothing specific, and invokes methods on the directory. Release note: None
- Loading branch information
1 parent
6123205
commit fd3c5be
Showing
12 changed files
with
355 additions
and
78 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test") | ||
|
|
||
| go_library( | ||
| name = "balancer", | ||
| srcs = ["balancer.go"], | ||
| importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/balancer", | ||
| visibility = ["//visibility:public"], | ||
| deps = [ | ||
| "//pkg/ccl/sqlproxyccl/tenant", | ||
| "//pkg/roachpb", | ||
| ], | ||
| ) | ||
|
|
||
| go_test( | ||
| name = "balancer_test", | ||
| srcs = ["balancer_test.go"], | ||
| deps = [ | ||
| ":balancer", | ||
| "//pkg/ccl/sqlproxyccl/tenant/testutils", | ||
| "//pkg/roachpb", | ||
| "//pkg/util/leaktest", | ||
| "@com_github_stretchr_testify//require", | ||
| ], | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| // Copyright 2022 The Cockroach Authors. | ||
| // | ||
| // Licensed as a CockroachDB Enterprise file under the Cockroach Community | ||
| // License (the "License"); you may not use this file except in compliance with | ||
| // the License. You may obtain a copy of the License at | ||
| // | ||
| // https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt | ||
|
|
||
| package balancer | ||
|
|
||
| import ( | ||
| "context" | ||
|
|
||
| "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant" | ||
| "github.com/cockroachdb/cockroach/pkg/roachpb" | ||
| ) | ||
|
|
||
| // Balancer handles load balancing of SQL connections within the proxy. | ||
| type Balancer struct { | ||
| // directory corresponds to the tenant directory that stores information | ||
| // about SQL pods for each tenant. | ||
| directory tenant.Directory | ||
| } | ||
|
|
||
| // NewBalancer constructs a new Balancer instance that is responsible for | ||
| // load balancing SQL connections within the proxy. | ||
| func NewBalancer(directory tenant.Directory) *Balancer { | ||
| return &Balancer{directory: directory} | ||
| } | ||
|
|
||
| // ChoosePodAddr returns the IP address of one of this tenant's available pods. | ||
| // This applies a weighted load balancing algorithm when selecting a pod from a | ||
| // list. If the tenant is suspended (i.e. list is empty initially), and no pods | ||
| // are available, this will trigger resumption of the tenant, and return the IP | ||
| // address of the new pod. If the tenant cannot be resumed, this may return a | ||
| // GRPC FailedPrecondition error. | ||
| // | ||
| // If clusterName is non-empty, then a GRPC NotFound error is returned if no | ||
| // pods match the cluster's name. This can be used to ensure that the | ||
| // incoming SQL connection "knows" some additional information about the | ||
| // tenant, such as the name of the cluster, before being allowed to connect. | ||
| // Similarly, if the tenant does not exist (e.g. because it was deleted), | ||
| // this returns a GRPC NotFound error. | ||
| // | ||
| // Note that resuming a tenant requires directory server calls, so this can | ||
| // block for some time, until the resumption process is complete. | ||
| // | ||
| // TODO(jaylim-crl): Remap GRPC NotFound and FailedPrecondition errors to a | ||
| // concrete error type. These GRPC errors should be internal to the service | ||
| // directory. | ||
| func (b *Balancer) ChoosePodAddr( | ||
| ctx context.Context, tenantID roachpb.TenantID, clusterName string, | ||
| ) (string, error) { | ||
| // TODO(jaylim-crl): We currently choose the pod's address within | ||
| // selectTenantPod, which is called by EnsureTenantAddr. That logic has to | ||
| // be extracted out into the balancer. EnsureTenantAddr should only block | ||
| // until there is at least 1 pod, so the balancer would first call | ||
| // EnsureTenantAddr, followed by LookupTenantAddrs to retrieve a list of | ||
| // SQL pods. Finally, the balancer would apply the weighted load balancing | ||
| // algorithm on the list of pods. The tenant directory should just be | ||
| // responsible for reporting a list of SQL pods, and allowing callers to | ||
| // resume SQL pods (e.g. EnsureTenantAddr). | ||
| return b.directory.EnsureTenantAddr(ctx, tenantID, clusterName) | ||
| } | ||
|
|
||
| // ListPodAddrs returns the IP addresses of all available SQL pods for the given | ||
| // tenant. This returns a GRPC NotFound error if the tenant does not exist (e.g. | ||
| // it has not yet been created) or if it has not yet been fetched into the | ||
| // directory's cache (ListPodAddrs will never attempt to fetch it). Unlike | ||
| // ChoosePodAddr, if no SQL pods are available for the tenant, ListPodAddrs | ||
| // will return the empty set. | ||
| // | ||
| // TODO(jaylim-crl): Remap GRPC NotFound to a concrete error type. The GRPC | ||
| // NotFound error should be internal to the directory. This GRPC error should be | ||
| // internal to the service directory. | ||
| func (b *Balancer) ListPodAddrs(ctx context.Context, tenantID roachpb.TenantID) ([]string, error) { | ||
| return b.directory.LookupTenantAddrs(ctx, tenantID) | ||
| } | ||
|
|
||
| // ReportFailure reports a connection failure to the balancer. This will inform | ||
| // the balancer that the process at addr is unhealthy, so that it is less likely | ||
| // to be chosen. It will also refresh the directory's internal cache. | ||
| func (b *Balancer) ReportFailure( | ||
| ctx context.Context, tenantID roachpb.TenantID, addr string, | ||
| ) error { | ||
| return b.directory.ReportFailure(ctx, tenantID, addr) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,80 @@ | ||
| // Copyright 2022 The Cockroach Authors. | ||
| // | ||
| // Licensed as a CockroachDB Enterprise file under the Cockroach Community | ||
| // License (the "License"); you may not use this file except in compliance with | ||
| // the License. You may obtain a copy of the License at | ||
| // | ||
| // https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt | ||
|
|
||
| package balancer_test | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
|
|
||
| "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/balancer" | ||
| "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant/testutils" | ||
| "github.com/cockroachdb/cockroach/pkg/roachpb" | ||
| "github.com/cockroachdb/cockroach/pkg/util/leaktest" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| // TestBalancer currently tests that it calls the directory with the right | ||
| // arguments. Once we add more logic to the balancer, we'll update this | ||
| // accordingly. We don't handle the error case here since the tests in | ||
| // proxy_handler_test.go should handle that. | ||
| func TestBalancer(t *testing.T) { | ||
| defer leaktest.AfterTest(t)() | ||
| ctx := context.Background() | ||
| tenantID := roachpb.MakeTenantID(10) | ||
| const clusterName = "foo-bar" | ||
|
|
||
| testDir := &testutils.TestDirectory{ | ||
| EnsureTenantAddrFn: func( | ||
| fnCtx context.Context, | ||
| fnTenantID roachpb.TenantID, | ||
| fnClusterName string, | ||
| ) (string, error) { | ||
| require.Equal(t, ctx, fnCtx) | ||
| require.Equal(t, tenantID, fnTenantID) | ||
| require.Equal(t, clusterName, fnClusterName) | ||
| return "127.0.0.10:42", nil | ||
| }, | ||
| LookupTenantAddrsFn: func( | ||
| fnCtx context.Context, | ||
| fnTenantID roachpb.TenantID, | ||
| ) ([]string, error) { | ||
| require.Equal(t, ctx, fnCtx) | ||
| require.Equal(t, tenantID, fnTenantID) | ||
| return []string{"127.0.0.10:42"}, nil | ||
| }, | ||
| ReportFailureFn: func( | ||
| fnCtx context.Context, | ||
| fnTenantID roachpb.TenantID, | ||
| fnAddr string, | ||
| ) error { | ||
| require.Equal(t, ctx, fnCtx) | ||
| require.Equal(t, tenantID, fnTenantID) | ||
| require.Equal(t, "127.0.0.10:42", fnAddr) | ||
| return nil | ||
| }, | ||
| } | ||
| b := balancer.NewBalancer(testDir) | ||
|
|
||
| addr, err := b.ChoosePodAddr(ctx, tenantID, clusterName) | ||
| require.NoError(t, err) | ||
| require.Equal(t, "127.0.0.10:42", addr) | ||
|
|
||
| var addrs []string | ||
| addrs, err = b.ListPodAddrs(ctx, tenantID) | ||
| require.NoError(t, err) | ||
| require.Equal(t, []string{"127.0.0.10:42"}, addrs) | ||
|
|
||
| err = b.ReportFailure(ctx, tenantID, "127.0.0.10:42") | ||
| require.NoError(t, err) | ||
|
|
||
| ensureCount, lookupCount, reportCount := testDir.Counts() | ||
| require.Equal(t, 1, ensureCount) | ||
| require.Equal(t, 1, lookupCount) | ||
| require.Equal(t, 1, reportCount) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.