Skip to content

Commit

Permalink
ccl/sqlproxyccl: add pkg/ccl/sqlproxyccl/tenant/{servicedir,simpledir}
Browse files Browse the repository at this point in the history
Previously, all directory related files are located in a tenant package, which
can be confusing. This commit moves files into specific directories, and is
purely mechanical.

Release note: None
  • Loading branch information
jaylim-crl committed Mar 20, 2022
1 parent 56a3c0b commit 695f8e1
Show file tree
Hide file tree
Showing 19 changed files with 209 additions and 165 deletions.
3 changes: 2 additions & 1 deletion pkg/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,8 @@ ALL_TESTS = [
"//pkg/ccl/sqlproxyccl/denylist:denylist_test",
"//pkg/ccl/sqlproxyccl/idle:idle_test",
"//pkg/ccl/sqlproxyccl/interceptor:interceptor_test",
"//pkg/ccl/sqlproxyccl/tenant:tenant_test",
"//pkg/ccl/sqlproxyccl/tenant/servicedir:servicedir_test",
"//pkg/ccl/sqlproxyccl/tenant/simpledir:simpledir_test",
"//pkg/ccl/sqlproxyccl/throttler:throttler_test",
"//pkg/ccl/sqlproxyccl:sqlproxyccl_test",
"//pkg/ccl/storageccl/engineccl:engineccl_test",
Expand Down
3 changes: 3 additions & 0 deletions pkg/ccl/sqlproxyccl/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ go_library(
"//pkg/ccl/sqlproxyccl/idle",
"//pkg/ccl/sqlproxyccl/interceptor",
"//pkg/ccl/sqlproxyccl/tenant",
"//pkg/ccl/sqlproxyccl/tenant/servicedir",
"//pkg/ccl/sqlproxyccl/tenant/simpledir",
"//pkg/ccl/sqlproxyccl/throttler",
"//pkg/roachpb",
"//pkg/security/certmgr",
Expand Down Expand Up @@ -71,6 +73,7 @@ go_test(
"//pkg/ccl/sqlproxyccl/denylist",
"//pkg/ccl/sqlproxyccl/interceptor",
"//pkg/ccl/sqlproxyccl/tenant",
"//pkg/ccl/sqlproxyccl/tenant/simpledir",
"//pkg/ccl/sqlproxyccl/tenantdirsvr",
"//pkg/ccl/sqlproxyccl/throttler",
"//pkg/ccl/utilccl",
Expand Down
18 changes: 10 additions & 8 deletions pkg/ccl/sqlproxyccl/proxy_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ import (
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/denylist"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/idle"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant/servicedir"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant/simpledir"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/throttler"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/security/certmgr"
Expand Down Expand Up @@ -184,22 +186,22 @@ func newProxyHandler(
// and the pod watcher. When a pod enters the DRAINING state, the pod
// watcher will set the idle monitor to detect connections without
// activity and terminate them.
var dirOpts []tenant.DirOption
var dirOpts []servicedir.DirOption
if options.DrainTimeout != 0 {
handler.idleMonitor = idle.NewMonitor(ctx, options.DrainTimeout)

podWatcher := make(chan *tenant.Pod)
podWatcher := make(chan *servicedir.Pod)
go handler.startPodWatcher(ctx, podWatcher)
dirOpts = append(dirOpts, tenant.PodWatcher(podWatcher))
dirOpts = append(dirOpts, servicedir.PodWatcher(podWatcher))
}

client := tenant.NewDirectoryClient(conn)
handler.directory, err = tenant.NewServiceDirectory(ctx, stopper, client, dirOpts...)
client := servicedir.NewDirectoryClient(conn)
handler.directory, err = servicedir.NewServiceDirectory(ctx, stopper, client, dirOpts...)
if err != nil {
return nil, err
}
} else {
handler.directory = tenant.NewSimpleDirectory(handler.RoutingRule)
handler.directory = simpledir.NewSimpleDirectory(handler.RoutingRule)
}

return &handler, nil
Expand Down Expand Up @@ -388,13 +390,13 @@ func (handler *proxyHandler) handle(ctx context.Context, incomingConn *proxyConn
// are subject to an idle timeout that closes them after a short period of
// inactivity. If a pod transitions back to the RUNNING state or to the DELETING
// state, then the idle timeout needs to be cleared.
func (handler *proxyHandler) startPodWatcher(ctx context.Context, podWatcher chan *tenant.Pod) {
func (handler *proxyHandler) startPodWatcher(ctx context.Context, podWatcher chan *servicedir.Pod) {
for {
select {
case <-ctx.Done():
return
case pod := <-podWatcher:
if pod.State == tenant.DRAINING {
if pod.State == servicedir.DRAINING {
handler.idleMonitor.SetIdleChecks(pod.Addr)
} else {
// Clear idle checks either for RUNNING or DELETING.
Expand Down
5 changes: 3 additions & 2 deletions pkg/ccl/sqlproxyccl/proxy_handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/ccl/kvccl/kvtenantccl"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/denylist"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant/simpledir"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenantdirsvr"
"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/throttler"
"github.com/cockroachdb/cockroach/pkg/roachpb"
Expand Down Expand Up @@ -105,7 +106,7 @@ func TestBackendDownRetry(t *testing.T) {

callCount := 0
defer testutils.TestingHook(
&tenant.ResolveTCPAddr,
&simpledir.ResolveTCPAddr,
func(network, addr string) (*net.TCPAddr, error) {
callCount++
if callCount >= 3 {
Expand Down Expand Up @@ -659,7 +660,7 @@ func TestDirectoryConnect(t *testing.T) {
proxy, addr := newProxyServer(ctx, t, srv.Stopper(), opts)

t.Run("tenant not found", func(t *testing.T) {
defer testutils.TestingHook(&tenant.ResolveTCPAddr,
defer testutils.TestingHook(&simpledir.ResolveTCPAddr,
func(network, addr string) (*net.TCPAddr, error) {
// Expect fallback.
require.Equal(t, srv.ServingSQLAddr(), addr)
Expand Down
78 changes: 3 additions & 75 deletions pkg/ccl/sqlproxyccl/tenant/BUILD.bazel
Original file line number Diff line number Diff line change
@@ -1,82 +1,10 @@
load("@bazel_gomock//:gomock.bzl", "gomock")
load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
load("@rules_proto//proto:defs.bzl", "proto_library")

proto_library(
name = "tenant_proto",
srcs = ["directory.proto"],
strip_import_prefix = "/pkg",
visibility = ["//visibility:public"],
deps = ["@com_github_gogo_protobuf//gogoproto:gogo_proto"],
)

go_proto_library(
name = "tenant_go_proto",
compilers = ["//pkg/cmd/protoc-gen-gogoroach:protoc-gen-gogoroach_grpc_compiler"],
importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant",
proto = ":tenant_proto",
visibility = ["//visibility:public"],
deps = ["@com_github_gogo_protobuf//gogoproto"],
)
load("@io_bazel_rules_go//go:def.bzl", "go_library")

go_library(
name = "tenant",
srcs = [
"entry.go",
"pod.go",
"resolver.go",
"service_directory.go",
"simple_directory.go",
],
embed = [":tenant_go_proto"],
srcs = ["resolver.go"],
importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant",
visibility = ["//visibility:public"],
deps = [
"//pkg/roachpb",
"//pkg/util/grpcutil",
"//pkg/util/log",
"//pkg/util/randutil",
"//pkg/util/stop",
"//pkg/util/syncutil",
"//pkg/util/timeutil",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
],
)

go_test(
name = "tenant_test",
size = "large",
srcs = [
"main_test.go",
"pod_test.go",
"service_directory_test.go",
"simple_directory_test.go",
],
embed = [":tenant"],
deps = [
"//pkg/base",
"//pkg/ccl",
"//pkg/ccl/kvccl/kvtenantccl",
"//pkg/ccl/sqlproxyccl/tenantdirsvr",
"//pkg/ccl/utilccl",
"//pkg/roachpb",
"//pkg/security",
"//pkg/security/securitytest",
"//pkg/server",
"//pkg/sql",
"//pkg/testutils",
"//pkg/testutils/serverutils",
"//pkg/testutils/skip",
"//pkg/testutils/testcluster",
"//pkg/util/leaktest",
"//pkg/util/log",
"//pkg/util/randutil",
"//pkg/util/stop",
"@com_github_stretchr_testify//require",
"@org_golang_google_grpc//:go_default_library",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
],
deps = ["//pkg/roachpb"],
)
78 changes: 78 additions & 0 deletions pkg/ccl/sqlproxyccl/tenant/servicedir/BUILD.bazel
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
load("@rules_proto//proto:defs.bzl", "proto_library")
load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")

proto_library(
name = "servicedir_proto",
srcs = ["directory.proto"],
strip_import_prefix = "/pkg",
visibility = ["//visibility:public"],
deps = ["@com_github_gogo_protobuf//gogoproto:gogo_proto"],
)

go_proto_library(
name = "servicedir_go_proto",
compilers = ["//pkg/cmd/protoc-gen-gogoroach:protoc-gen-gogoroach_grpc_compiler"],
importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant/servicedir",
proto = ":servicedir_proto",
visibility = ["//visibility:public"],
deps = ["@com_github_gogo_protobuf//gogoproto"],
)

go_library(
name = "servicedir",
srcs = [
"entry.go",
"pod.go",
"service_directory.go",
],
embed = [":servicedir_go_proto"],
importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant/servicedir",
visibility = ["//visibility:public"],
deps = [
"//pkg/ccl/sqlproxyccl/tenant",
"//pkg/roachpb",
"//pkg/util/grpcutil",
"//pkg/util/log",
"//pkg/util/randutil",
"//pkg/util/stop",
"//pkg/util/syncutil",
"//pkg/util/timeutil",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
],
)

go_test(
name = "servicedir_test",
srcs = [
"main_test.go",
"pod_test.go",
"service_directory_test.go",
],
embed = [":servicedir"],
deps = [
"//pkg/base",
"//pkg/ccl",
"//pkg/ccl/kvccl/kvtenantccl",
"//pkg/ccl/sqlproxyccl/tenant",
"//pkg/ccl/sqlproxyccl/tenantdirsvr",
"//pkg/ccl/utilccl",
"//pkg/roachpb",
"//pkg/security",
"//pkg/security/securitytest",
"//pkg/server",
"//pkg/sql",
"//pkg/testutils/serverutils",
"//pkg/testutils/skip",
"//pkg/testutils/testcluster",
"//pkg/util/leaktest",
"//pkg/util/log",
"//pkg/util/randutil",
"//pkg/util/stop",
"@com_github_stretchr_testify//require",
"@org_golang_google_grpc//:go_default_library",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
],
)
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright 2021 The Cockroach Authors.
// 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
Expand All @@ -7,8 +7,8 @@
// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt

syntax = "proto3";
package cockroach.ccl.sqlproxyccl.tenant;
option go_package="tenant";
package cockroach.ccl.sqlproxyccl.tenant.servicedir;
option go_package="servicedir";

import "gogoproto/gogo.proto";

Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
// Copyright 2021 The Cockroach Authors.
// 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 tenant
package servicedir

import (
"context"
Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
// Copyright 2018 The Cockroach Authors.
// 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 tenant_test
package servicedir_test

import (
"os"
Expand All @@ -31,4 +31,4 @@ func TestMain(m *testing.M) {
os.Exit(m.Run())
}

//go:generate ../../../util/leaktest/add-leaktest.sh *_test.go
//go:generate ../../../../util/leaktest/add-leaktest.sh *_test.go
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
// Copyright 2021 The Cockroach Authors.
// 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 tenant
package servicedir

// selectTenantPod selects a tenant pod from the given list to received
// incoming traffic. Pods are weighted by their reported CPU load. rand must be
Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
// Copyright 2021 The Cockroach Authors.
// 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 tenant
package servicedir

import (
"math/rand"
Expand Down
Original file line number Diff line number Diff line change
@@ -1,19 +1,20 @@
// Copyright 2021 The Cockroach Authors.
// 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 tenant
package servicedir

import (
"context"
"io"
"sync"
"time"

"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
"github.com/cockroachdb/cockroach/pkg/util/log"
Expand Down Expand Up @@ -106,7 +107,7 @@ type serviceDirectory struct {
// needed.
func NewServiceDirectory(
ctx context.Context, stopper *stop.Stopper, client DirectoryClient, opts ...DirOption,
) (Resolver, error) {
) (tenant.Resolver, error) {
dir := &serviceDirectory{client: client, stopper: stopper}

dir.mut.tenants = make(map[roachpb.TenantID]*tenantEntry)
Expand Down
Loading

0 comments on commit 695f8e1

Please sign in to comment.