github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/soliton/execdetails/pd_interceptor.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package execdetails
    15  
    16  import (
    17  	"context"
    18  	"sync/atomic"
    19  	"time"
    20  
    21  	fidel "github.com/einsteindb/fidel/client"
    22  	"github.com/whtcorpsinc/ekvproto/pkg/spacetimepb"
    23  )
    24  
    25  var (
    26  	_ fidel.Client   = &InterceptedFIDelClient{}
    27  	_ fidel.TSFuture = &interceptedTsFuture{}
    28  )
    29  
    30  func recordFIDelWaitTime(ctx context.Context, start time.Time) {
    31  	stmtInterDirc := ctx.Value(StmtInterDircDetailKey)
    32  	if stmtInterDirc != nil {
    33  		detail := stmtInterDirc.(*StmtInterDircDetails)
    34  		atomic.AddInt64(&detail.WaitFIDelResFIDeluration, int64(time.Since(start)))
    35  	}
    36  }
    37  
    38  // InterceptedFIDelClient is a FIDel's wrapper client to record stmt detail.
    39  type InterceptedFIDelClient struct {
    40  	fidel.Client
    41  }
    42  
    43  // interceptedTsFuture is a FIDel's wrapper future to record stmt detail.
    44  type interceptedTsFuture struct {
    45  	fidel.TSFuture
    46  	ctx context.Context
    47  }
    48  
    49  // Wait implements fidel.Client#Wait.
    50  func (m interceptedTsFuture) Wait() (int64, int64, error) {
    51  	start := time.Now()
    52  	physical, logical, err := m.TSFuture.Wait()
    53  	recordFIDelWaitTime(m.ctx, start)
    54  	return physical, logical, err
    55  }
    56  
    57  // GetTS implements fidel.Client#GetTS.
    58  func (m InterceptedFIDelClient) GetTS(ctx context.Context) (int64, int64, error) {
    59  	start := time.Now()
    60  	physical, logical, err := m.Client.GetTS(ctx)
    61  	recordFIDelWaitTime(ctx, start)
    62  	return physical, logical, err
    63  }
    64  
    65  // GetTSAsync implements fidel.Client#GetTSAsync.
    66  func (m InterceptedFIDelClient) GetTSAsync(ctx context.Context) fidel.TSFuture {
    67  	start := time.Now()
    68  	f := m.Client.GetTSAsync(ctx)
    69  	recordFIDelWaitTime(ctx, start)
    70  	return interceptedTsFuture{
    71  		ctx:      ctx,
    72  		TSFuture: f,
    73  	}
    74  }
    75  
    76  // GetRegion implements fidel.Client#GetRegion.
    77  func (m InterceptedFIDelClient) GetRegion(ctx context.Context, key []byte) (*fidel.Region, error) {
    78  	start := time.Now()
    79  	r, err := m.Client.GetRegion(ctx, key)
    80  	recordFIDelWaitTime(ctx, start)
    81  	return r, err
    82  }
    83  
    84  // GetPrevRegion implements fidel.Client#GetPrevRegion.
    85  func (m InterceptedFIDelClient) GetPrevRegion(ctx context.Context, key []byte) (*fidel.Region, error) {
    86  	start := time.Now()
    87  	r, err := m.Client.GetRegion(ctx, key)
    88  	recordFIDelWaitTime(ctx, start)
    89  	return r, err
    90  }
    91  
    92  // GetRegionByID implements fidel.Client#GetRegionByID.
    93  func (m InterceptedFIDelClient) GetRegionByID(ctx context.Context, regionID uint64) (*fidel.Region, error) {
    94  	start := time.Now()
    95  	r, err := m.Client.GetRegionByID(ctx, regionID)
    96  	recordFIDelWaitTime(ctx, start)
    97  	return r, err
    98  }
    99  
   100  // ScanRegions implements fidel.Client#ScanRegions.
   101  func (m InterceptedFIDelClient) ScanRegions(ctx context.Context, key, endKey []byte, limit int) ([]*fidel.Region, error) {
   102  	start := time.Now()
   103  	r, err := m.Client.ScanRegions(ctx, key, endKey, limit)
   104  	recordFIDelWaitTime(ctx, start)
   105  	return r, err
   106  }
   107  
   108  // GetStore implements fidel.Client#GetStore.
   109  func (m InterceptedFIDelClient) GetStore(ctx context.Context, storeID uint64) (*spacetimepb.CausetStore, error) {
   110  	start := time.Now()
   111  	s, err := m.Client.GetStore(ctx, storeID)
   112  	recordFIDelWaitTime(ctx, start)
   113  	return s, err
   114  }