github.com/KinWaiYuen/client-go/v2@v2.5.4/util/pd_interceptor.go (about)

     1  // Copyright 2021 TiKV Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // NOTE: The code in this file is based on code from the
    16  // TiDB project, licensed under the Apache License v 2.0
    17  //
    18  // https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/util/pd_interceptor.go
    19  //
    20  
    21  // Copyright 2020 PingCAP, Inc.
    22  //
    23  // Licensed under the Apache License, Version 2.0 (the "License");
    24  // you may not use this file except in compliance with the License.
    25  // You may obtain a copy of the License at
    26  //
    27  //     http://www.apache.org/licenses/LICENSE-2.0
    28  //
    29  // Unless required by applicable law or agreed to in writing, software
    30  // distributed under the License is distributed on an "AS IS" BASIS,
    31  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    32  // See the License for the specific language governing permissions and
    33  // limitations under the License.
    34  
    35  package util
    36  
    37  import (
    38  	"context"
    39  	"sync/atomic"
    40  	"time"
    41  
    42  	"github.com/pingcap/kvproto/pkg/metapb"
    43  	pd "github.com/tikv/pd/client"
    44  )
    45  
    46  var (
    47  	_ pd.Client   = &InterceptedPDClient{}
    48  	_ pd.TSFuture = &interceptedTsFuture{}
    49  )
    50  
    51  func recordPDWaitTime(ctx context.Context, start time.Time) {
    52  	stmtExec := ctx.Value(ExecDetailsKey)
    53  	if stmtExec != nil {
    54  		detail := stmtExec.(*ExecDetails)
    55  		atomic.AddInt64(&detail.WaitPDRespDuration, int64(time.Since(start)))
    56  	}
    57  }
    58  
    59  // InterceptedPDClient is a PD's wrapper client to record stmt detail.
    60  type InterceptedPDClient struct {
    61  	pd.Client
    62  }
    63  
    64  // interceptedTsFuture is a PD's wrapper future to record stmt detail.
    65  type interceptedTsFuture struct {
    66  	pd.TSFuture
    67  	ctx context.Context
    68  }
    69  
    70  // Wait implements pd.Client#Wait.
    71  func (m interceptedTsFuture) Wait() (int64, int64, error) {
    72  	start := time.Now()
    73  	physical, logical, err := m.TSFuture.Wait()
    74  	recordPDWaitTime(m.ctx, start)
    75  	return physical, logical, err
    76  }
    77  
    78  // GetTS implements pd.Client#GetTS.
    79  func (m InterceptedPDClient) GetTS(ctx context.Context) (int64, int64, error) {
    80  	start := time.Now()
    81  	physical, logical, err := m.Client.GetTS(ctx)
    82  	recordPDWaitTime(ctx, start)
    83  	return physical, logical, err
    84  }
    85  
    86  // GetTSAsync implements pd.Client#GetTSAsync.
    87  func (m InterceptedPDClient) GetTSAsync(ctx context.Context) pd.TSFuture {
    88  	start := time.Now()
    89  	f := m.Client.GetTSAsync(ctx)
    90  	recordPDWaitTime(ctx, start)
    91  	return interceptedTsFuture{
    92  		ctx:      ctx,
    93  		TSFuture: f,
    94  	}
    95  }
    96  
    97  // GetRegion implements pd.Client#GetRegion.
    98  func (m InterceptedPDClient) GetRegion(ctx context.Context, key []byte) (*pd.Region, error) {
    99  	start := time.Now()
   100  	r, err := m.Client.GetRegion(ctx, key)
   101  	recordPDWaitTime(ctx, start)
   102  	return r, err
   103  }
   104  
   105  // GetPrevRegion implements pd.Client#GetPrevRegion.
   106  func (m InterceptedPDClient) GetPrevRegion(ctx context.Context, key []byte) (*pd.Region, error) {
   107  	start := time.Now()
   108  	r, err := m.Client.GetPrevRegion(ctx, key)
   109  	recordPDWaitTime(ctx, start)
   110  	return r, err
   111  }
   112  
   113  // GetRegionByID implements pd.Client#GetRegionByID.
   114  func (m InterceptedPDClient) GetRegionByID(ctx context.Context, regionID uint64) (*pd.Region, error) {
   115  	start := time.Now()
   116  	r, err := m.Client.GetRegionByID(ctx, regionID)
   117  	recordPDWaitTime(ctx, start)
   118  	return r, err
   119  }
   120  
   121  // ScanRegions implements pd.Client#ScanRegions.
   122  func (m InterceptedPDClient) ScanRegions(ctx context.Context, key, endKey []byte, limit int) ([]*pd.Region, error) {
   123  	start := time.Now()
   124  	r, err := m.Client.ScanRegions(ctx, key, endKey, limit)
   125  	recordPDWaitTime(ctx, start)
   126  	return r, err
   127  }
   128  
   129  // GetStore implements pd.Client#GetStore.
   130  func (m InterceptedPDClient) GetStore(ctx context.Context, storeID uint64) (*metapb.Store, error) {
   131  	start := time.Now()
   132  	s, err := m.Client.GetStore(ctx, storeID)
   133  	recordPDWaitTime(ctx, start)
   134  	return s, err
   135  }