Skip to content

Commit

Permalink
refactor(): update client-map's controlling access
Browse files Browse the repository at this point in the history
  • Loading branch information
shhdgit committed Mar 23, 2021
1 parent a6a5c84 commit 5b6cd66
Show file tree
Hide file tree
Showing 6 changed files with 121 additions and 98 deletions.
115 changes: 115 additions & 0 deletions pkg/apiserver/profiling/clientmap.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
// Copyright 2021 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package profiling

import (
"fmt"
"net/http"
"time"

"github.com/pingcap/tidb-dashboard/pkg/apiserver/model"
"github.com/pingcap/tidb-dashboard/pkg/apiserver/profiling/fetcher"
"github.com/pingcap/tidb-dashboard/pkg/config"
"github.com/pingcap/tidb-dashboard/pkg/pd"
"github.com/pingcap/tidb-dashboard/pkg/tidb"
"github.com/pingcap/tidb-dashboard/pkg/tiflash"
"github.com/pingcap/tidb-dashboard/pkg/tikv"
)

const (
maxProfilingTimeout = time.Minute * 5
)

type clientMap map[model.NodeKind]fetcher.Client

func (fm *clientMap) Get(kind model.NodeKind) (fetcher.Client, error) {
f, ok := (*fm)[kind]
if !ok {
return nil, fmt.Errorf("unsupported target %s", kind)
}
return f, nil
}

func newClientMap(
tikvHTTPClient *tikv.Client,
tiflashHTTPClient *tiflash.Client,
tidbHTTPClient *tidb.Client,
pdHTTPClient *pd.Client,
config *config.Config,
) *clientMap {
return &clientMap{
model.NodeKindTiKV: &tikvClient{
client: tikvHTTPClient,
},
model.NodeKindTiFlash: &tiflashClient{
client: tiflashHTTPClient,
},
model.NodeKindTiDB: &tidbClient{
client: tidbHTTPClient,
},
model.NodeKindPD: &pdClient{
client: pdHTTPClient,
statusAPIHTTPScheme: config.GetClusterHTTPScheme(),
},
}
}

// tikv
var _ fetcher.Client = (*tikvClient)(nil)

type tikvClient struct {
client *tikv.Client
}

func (f *tikvClient) Fetch(op *fetcher.ClientFetchOptions) ([]byte, error) {
return f.client.WithTimeout(maxProfilingTimeout).SendGetRequest(op.IP, op.Port, op.Path)
}

// tiflash
var _ fetcher.Client = (*tiflashClient)(nil)

type tiflashClient struct {
client *tiflash.Client
}

func (f *tiflashClient) Fetch(op *fetcher.ClientFetchOptions) ([]byte, error) {
return f.client.WithTimeout(maxProfilingTimeout).SendGetRequest(op.IP, op.Port, op.Path)
}

// tidb
var _ fetcher.Client = (*tidbClient)(nil)

type tidbClient struct {
client *tidb.Client
}

func (f *tidbClient) Fetch(op *fetcher.ClientFetchOptions) ([]byte, error) {
return f.client.WithStatusAPIAddress(op.IP, op.Port).WithStatusAPITimeout(maxProfilingTimeout).SendGetRequest(op.Path)
}

// pd
var _ fetcher.Client = (*pdClient)(nil)

type pdClient struct {
client *pd.Client
statusAPIHTTPScheme string
}

func (f *pdClient) Fetch(op *fetcher.ClientFetchOptions) ([]byte, error) {
baseURL := fmt.Sprintf("%s://%s:%d", f.statusAPIHTTPScheme, op.IP, op.Port)
f.client.WithBeforeRequest(func(req *http.Request) {
req.Header.Add("PD-Allow-follower-handle", "true")
})
return f.client.WithTimeout(maxProfilingTimeout).WithBaseURL(baseURL).SendGetRequest(op.Path)
}
88 changes: 0 additions & 88 deletions pkg/apiserver/profiling/fetcher/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,6 @@

package fetcher

import (
"fmt"
"net/http"
"time"

"github.com/pingcap/tidb-dashboard/pkg/apiserver/model"
"github.com/pingcap/tidb-dashboard/pkg/config"
"github.com/pingcap/tidb-dashboard/pkg/pd"
"github.com/pingcap/tidb-dashboard/pkg/tidb"
"github.com/pingcap/tidb-dashboard/pkg/tiflash"
"github.com/pingcap/tidb-dashboard/pkg/tikv"
)

const (
maxProfilingTimeout = time.Minute * 5
)

type ClientFetchOptions struct {
IP string
Port int
Expand All @@ -39,74 +22,3 @@ type ClientFetchOptions struct {
type Client interface {
Fetch(op *ClientFetchOptions) ([]byte, error)
}

type ClientMap map[model.NodeKind]Client

func (fm *ClientMap) Get(kind model.NodeKind) (Client, error) {
f, ok := (*fm)[kind]
if !ok {
return nil, fmt.Errorf("unsupported target %s", kind)
}
return f, nil
}

func NewClientMap(
tikvHTTPClient *tikv.Client,
tiflashHTTPClient *tiflash.Client,
tidbHTTPClient *tidb.Client,
pdHTTPClient *pd.Client,
config *config.Config,
) *ClientMap {
return &ClientMap{
model.NodeKindTiKV: &tikvClient{
client: tikvHTTPClient,
},
model.NodeKindTiFlash: &tiflashClient{
client: tiflashHTTPClient,
},
model.NodeKindTiDB: &tidbClient{
client: tidbHTTPClient,
},
model.NodeKindPD: &pdClient{
client: pdHTTPClient,
statusAPIHTTPScheme: config.GetClusterHTTPScheme(),
},
}
}

type tikvClient struct {
client *tikv.Client
}

func (f *tikvClient) Fetch(op *ClientFetchOptions) ([]byte, error) {
return f.client.WithTimeout(maxProfilingTimeout).SendGetRequest(op.IP, op.Port, op.Path)
}

type tiflashClient struct {
client *tiflash.Client
}

func (f *tiflashClient) Fetch(op *ClientFetchOptions) ([]byte, error) {
return f.client.WithTimeout(maxProfilingTimeout).SendGetRequest(op.IP, op.Port, op.Path)
}

type tidbClient struct {
client *tidb.Client
}

func (f *tidbClient) Fetch(op *ClientFetchOptions) ([]byte, error) {
return f.client.WithStatusAPIAddress(op.IP, op.Port).WithStatusAPITimeout(maxProfilingTimeout).SendGetRequest(op.Path)
}

type pdClient struct {
client *pd.Client
statusAPIHTTPScheme string
}

func (f *pdClient) Fetch(op *ClientFetchOptions) ([]byte, error) {
baseURL := fmt.Sprintf("%s://%s:%d", f.statusAPIHTTPScheme, op.IP, op.Port)
f.client.WithBeforeRequest(func(req *http.Request) {
req.Header.Add("PD-Allow-follower-handle", "true")
})
return f.client.WithTimeout(maxProfilingTimeout).WithBaseURL(baseURL).SendGetRequest(op.Path)
}
5 changes: 2 additions & 3 deletions pkg/apiserver/profiling/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ import (
"time"

"github.com/pingcap/tidb-dashboard/pkg/apiserver/model"
"github.com/pingcap/tidb-dashboard/pkg/apiserver/profiling/fetcher"
"github.com/pingcap/tidb-dashboard/pkg/dbstore"
)

Expand Down Expand Up @@ -73,11 +72,11 @@ type Task struct {
ctx context.Context
cancel context.CancelFunc
taskGroup *TaskGroup
clientMap *fetcher.ClientMap
clientMap *clientMap
}

// NewTask creates a new profiling task.
func NewTask(ctx context.Context, taskGroup *TaskGroup, target model.RequestTargetNode, cm *fetcher.ClientMap) *Task {
func NewTask(ctx context.Context, taskGroup *TaskGroup, target model.RequestTargetNode, cm *clientMap) *Task {
ctx, cancel := context.WithCancel(ctx)
return &Task{
TaskModel: &TaskModel{
Expand Down
4 changes: 1 addition & 3 deletions pkg/apiserver/profiling/module.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,11 @@ package profiling

import (
"go.uber.org/fx"

"github.com/pingcap/tidb-dashboard/pkg/apiserver/profiling/fetcher"
)

var Module = fx.Options(
fx.Provide(
fetcher.NewClientMap,
newClientMap,
newService,
),
fx.Invoke(registerRouter),
Expand Down
2 changes: 1 addition & 1 deletion pkg/apiserver/profiling/profile.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (
"github.com/pingcap/tidb-dashboard/pkg/apiserver/profiling/fetcher"
)

func profileAndWriteSVG(ctx context.Context, cm *fetcher.ClientMap, target *model.RequestTargetNode, fileNameWithoutExt string, profileDurationSecs uint) (string, error) {
func profileAndWriteSVG(ctx context.Context, cm *clientMap, target *model.RequestTargetNode, fileNameWithoutExt string, profileDurationSecs uint) (string, error) {
c, err := cm.Get(target.Kind)
if err != nil {
return "", err
Expand Down
5 changes: 2 additions & 3 deletions pkg/apiserver/profiling/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ import (
"go.uber.org/zap"

"github.com/pingcap/tidb-dashboard/pkg/apiserver/model"
"github.com/pingcap/tidb-dashboard/pkg/apiserver/profiling/fetcher"
"github.com/pingcap/tidb-dashboard/pkg/config"
"github.com/pingcap/tidb-dashboard/pkg/dbstore"
)
Expand Down Expand Up @@ -63,10 +62,10 @@ type Service struct {
sessionCh chan *StartRequestSession
lastTaskGroup *TaskGroup
tasks sync.Map
clientMap *fetcher.ClientMap
clientMap *clientMap
}

func newService(lc fx.Lifecycle, p ServiceParams, cm *fetcher.ClientMap) (*Service, error) {
func newService(lc fx.Lifecycle, p ServiceParams, cm *clientMap) (*Service, error) {
if err := autoMigrate(p.LocalStore); err != nil {
return nil, err
}
Expand Down

0 comments on commit 5b6cd66

Please sign in to comment.