github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/milevadb-server/einsteindb/client_collapse.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 einsteindb provides tcp connection to ekvserver.
    15  package einsteindb
    16  
    17  import (
    18  	"context"
    19  	"strconv"
    20  	"time"
    21  
    22  	"github.com/whtcorpsinc/errors"
    23  	"github.com/whtcorpsinc/milevadb/causetstore/einsteindb/einsteindbrpc"
    24  	"golang.org/x/sync/singleflight"
    25  )
    26  
    27  var _ Client = reqDefCauslapse{}
    28  
    29  var resolveRegionSf singleflight.Group
    30  
    31  type reqDefCauslapse struct {
    32  	Client
    33  }
    34  
    35  func (r reqDefCauslapse) Close() error {
    36  	if r.Client == nil {
    37  		panic("client should not be nil")
    38  	}
    39  	return r.Client.Close()
    40  }
    41  
    42  func (r reqDefCauslapse) SendRequest(ctx context.Context, addr string, req *einsteindbrpc.Request, timeout time.Duration) (*einsteindbrpc.Response, error) {
    43  	if r.Client == nil {
    44  		panic("client should not be nil")
    45  	}
    46  	if canDefCauslapse, resp, err := r.tryDefCauslapseRequest(ctx, addr, req, timeout); canDefCauslapse {
    47  		return resp, err
    48  	}
    49  	return r.Client.SendRequest(ctx, addr, req, timeout)
    50  }
    51  
    52  func (r reqDefCauslapse) tryDefCauslapseRequest(ctx context.Context, addr string, req *einsteindbrpc.Request, timeout time.Duration) (canDefCauslapse bool, resp *einsteindbrpc.Response, err error) {
    53  	switch req.Type {
    54  	case einsteindbrpc.CmdResolveLock:
    55  		resolveLock := req.ResolveLock()
    56  		if len(resolveLock.Keys) > 0 {
    57  			// can not collapse resolve dagger lite
    58  			return
    59  		}
    60  		if len(resolveLock.TxnInfos) > 0 {
    61  			// can not collapse batch resolve locks which is only used by GC worker.
    62  			return
    63  		}
    64  		canDefCauslapse = true
    65  		key := strconv.FormatUint(resolveLock.Context.RegionId, 10) + "-" + strconv.FormatUint(resolveLock.StartVersion, 10)
    66  		resp, err = r.collapse(ctx, key, &resolveRegionSf, addr, req, timeout)
    67  		return
    68  	default:
    69  		// now we only support collapse resolve dagger.
    70  		return
    71  	}
    72  }
    73  
    74  func (r reqDefCauslapse) collapse(ctx context.Context, key string, sf *singleflight.Group,
    75  	addr string, req *einsteindbrpc.Request, timeout time.Duration) (resp *einsteindbrpc.Response, err error) {
    76  	rsC := sf.DoChan(key, func() (interface{}, error) {
    77  		return r.Client.SendRequest(context.Background(), addr, req, readTimeoutShort) // use resolveLock timeout.
    78  	})
    79  	timer := time.NewTimer(timeout)
    80  	defer timer.Stop()
    81  	select {
    82  	case <-ctx.Done():
    83  		err = errors.Trace(ctx.Err())
    84  		return
    85  	case <-timer.C:
    86  		err = errors.Trace(context.DeadlineExceeded)
    87  		return
    88  	case rs := <-rsC:
    89  		if rs.Err != nil {
    90  			err = errors.Trace(rs.Err)
    91  			return
    92  		}
    93  		resp = rs.Val.(*einsteindbrpc.Response)
    94  		return
    95  	}
    96  }