github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/causetstore/entangledstore/dag_poset_handler/dap_poset_handler.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 cophandler
    15  
    16  import (
    17  	"bytes"
    18  	"fmt"
    19  	"time"
    20  
    21  	"github.com/golang/protobuf/proto"
    22  	"github.com/ngaut/entangledstore/einsteindb/dbreader"
    23  	"github.com/ngaut/entangledstore/lockstore"
    24  	"github.com/whtcorpsinc/BerolinaSQL/allegrosql"
    25  	"github.com/whtcorpsinc/BerolinaSQL/perceptron"
    26  	"github.com/whtcorpsinc/BerolinaSQL/terror"
    27  	"github.com/whtcorpsinc/ekvproto/pkg/ekvrpcpb"
    28  	"github.com/whtcorpsinc/ekvproto/pkg/interlock"
    29  	"github.com/whtcorpsinc/errors"
    30  	"github.com/whtcorpsinc/fidelpb/go-fidelpb"
    31  	"github.com/whtcorpsinc/milevadb/blockcodec"
    32  	"github.com/whtcorpsinc/milevadb/ekv"
    33  	"github.com/whtcorpsinc/milevadb/memex"
    34  	"github.com/whtcorpsinc/milevadb/memex/aggregation"
    35  	"github.com/whtcorpsinc/milevadb/soliton/chunk"
    36  	"github.com/whtcorpsinc/milevadb/soliton/codec"
    37  	"github.com/whtcorpsinc/milevadb/soliton/collate"
    38  	"github.com/whtcorpsinc/milevadb/soliton/rowcodec"
    39  	"github.com/whtcorpsinc/milevadb/stochastikctx/stmtctx"
    40  	"github.com/whtcorpsinc/milevadb/types"
    41  )
    42  
    43  // HandleCopRequest handles interlock request.
    44  func HandleCopRequest(dbReader *dbreader.DBReader, lockStore *lockstore.MemStore, req *interlock.Request) *interlock.Response {
    45  	switch req.Tp {
    46  	case ekv.ReqTypePosetDag:
    47  		return handleCoFIDelAGRequest(dbReader, lockStore, req)
    48  	case ekv.ReqTypeAnalyze:
    49  		return handleCopAnalyzeRequest(dbReader, req)
    50  	case ekv.ReqTypeChecksum:
    51  		return handleCopChecksumRequest(dbReader, req)
    52  	}
    53  	return &interlock.Response{OtherError: fmt.Sprintf("unsupported request type %d", req.GetTp())}
    54  }
    55  
    56  type posetPosetDagContext struct {
    57  	*evalContext
    58  	dbReader         *dbreader.DBReader
    59  	lockStore        *lockstore.MemStore
    60  	resolvedLocks    []uint64
    61  	posetPosetDagReq *fidelpb.PosetDagRequest
    62  	keyRanges        []*interlock.KeyRange
    63  	startTS          uint64
    64  }
    65  
    66  // handleCoFIDelAGRequest handles interlock PosetDag request.
    67  func handleCoFIDelAGRequest(dbReader *dbreader.DBReader, lockStore *lockstore.MemStore, req *interlock.Request) *interlock.Response {
    68  	startTime := time.Now()
    69  	resp := &interlock.Response{}
    70  	posetPosetDagCtx, posetPosetDagReq, err := buildPosetDag(dbReader, lockStore, req)
    71  	if err != nil {
    72  		resp.OtherError = err.Error()
    73  		return resp
    74  	}
    75  	closureInterDirc, err := buildClosureInterlockingDirectorate(posetPosetDagCtx, posetPosetDagReq)
    76  	if err != nil {
    77  		return buildResp(nil, nil, posetPosetDagReq, err, posetPosetDagCtx.sc.GetWarnings(), time.Since(startTime))
    78  	}
    79  	chunks, err := closureInterDirc.execute()
    80  	return buildResp(chunks, closureInterDirc.counts, posetPosetDagReq, err, posetPosetDagCtx.sc.GetWarnings(), time.Since(startTime))
    81  }
    82  
    83  func buildPosetDag(reader *dbreader.DBReader, lockStore *lockstore.MemStore, req *interlock.Request) (*posetPosetDagContext, *fidelpb.PosetDagRequest, error) {
    84  	if len(req.Ranges) == 0 {
    85  		return nil, nil, errors.New("request range is null")
    86  	}
    87  	if req.GetTp() != ekv.ReqTypePosetDag {
    88  		return nil, nil, errors.Errorf("unsupported request type %d", req.GetTp())
    89  	}
    90  
    91  	posetPosetDagReq := new(fidelpb.PosetDagRequest)
    92  	err := proto.Unmarshal(req.Data, posetPosetDagReq)
    93  	if err != nil {
    94  		return nil, nil, errors.Trace(err)
    95  	}
    96  	sc := flagsToStatementContext(posetPosetDagReq.Flags)
    97  	sc.TimeZone = time.FixedZone("UTC", int(posetPosetDagReq.TimeZoneOffset))
    98  	ctx := &posetPosetDagContext{
    99  		evalContext:      &evalContext{sc: sc},
   100  		dbReader:         reader,
   101  		lockStore:        lockStore,
   102  		posetPosetDagReq: posetPosetDagReq,
   103  		keyRanges:        req.Ranges,
   104  		startTS:          req.StartTs,
   105  		resolvedLocks:    req.Context.ResolvedLocks,
   106  	}
   107  	scanInterDirc := posetPosetDagReq.InterlockingDirectorates[0]
   108  	if scanInterDirc.Tp == fidelpb.InterDircType_TypeTableScan {
   109  		ctx.setDeferredCausetInfo(scanInterDirc.TblScan.DeferredCausets)
   110  		ctx.primaryDefCauss = scanInterDirc.TblScan.PrimaryDeferredCausetIds
   111  	} else {
   112  		ctx.setDeferredCausetInfo(scanInterDirc.IdxScan.DeferredCausets)
   113  	}
   114  	return ctx, posetPosetDagReq, err
   115  }
   116  
   117  func getAggInfo(ctx *posetPosetDagContext, pbAgg *fidelpb.Aggregation) ([]aggregation.Aggregation, []memex.Expression, error) {
   118  	length := len(pbAgg.AggFunc)
   119  	aggs := make([]aggregation.Aggregation, 0, length)
   120  	var err error
   121  	for _, expr := range pbAgg.AggFunc {
   122  		var aggExpr aggregation.Aggregation
   123  		aggExpr, err = aggregation.NewDistAggFunc(expr, ctx.fieldTps, ctx.sc)
   124  		if err != nil {
   125  			return nil, nil, errors.Trace(err)
   126  		}
   127  		aggs = append(aggs, aggExpr)
   128  	}
   129  	groupBys, err := convertToExprs(ctx.sc, ctx.fieldTps, pbAgg.GetGroupBy())
   130  	if err != nil {
   131  		return nil, nil, errors.Trace(err)
   132  	}
   133  
   134  	return aggs, groupBys, nil
   135  }
   136  
   137  func getTopNInfo(ctx *evalContext, topN *fidelpb.TopN) (heap *topNHeap, conds []memex.Expression, err error) {
   138  	pbConds := make([]*fidelpb.Expr, len(topN.OrderBy))
   139  	for i, item := range topN.OrderBy {
   140  		pbConds[i] = item.Expr
   141  	}
   142  	heap = &topNHeap{
   143  		totalCount: int(topN.Limit),
   144  		topNSorter: topNSorter{
   145  			orderByItems: topN.OrderBy,
   146  			sc:           ctx.sc,
   147  		},
   148  	}
   149  	if conds, err = convertToExprs(ctx.sc, ctx.fieldTps, pbConds); err != nil {
   150  		return nil, nil, errors.Trace(err)
   151  	}
   152  
   153  	return heap, conds, nil
   154  }
   155  
   156  type evalContext struct {
   157  	colIDs          map[int64]int
   158  	columnInfos     []*fidelpb.DeferredCausetInfo
   159  	fieldTps        []*types.FieldType
   160  	primaryDefCauss []int64
   161  	sc              *stmtctx.StatementContext
   162  }
   163  
   164  func (e *evalContext) setDeferredCausetInfo(defcaus []*fidelpb.DeferredCausetInfo) {
   165  	e.columnInfos = make([]*fidelpb.DeferredCausetInfo, len(defcaus))
   166  	copy(e.columnInfos, defcaus)
   167  
   168  	e.colIDs = make(map[int64]int, len(e.columnInfos))
   169  	e.fieldTps = make([]*types.FieldType, 0, len(e.columnInfos))
   170  	for i, col := range e.columnInfos {
   171  		ft := fieldTypeFromPBDeferredCauset(col)
   172  		e.fieldTps = append(e.fieldTps, ft)
   173  		e.colIDs[col.GetDeferredCausetId()] = i
   174  	}
   175  }
   176  
   177  func (e *evalContext) newRowCausetDecoder() (*rowcodec.ChunkCausetDecoder, error) {
   178  	var (
   179  		pkDefCauss []int64
   180  		defcaus    = make([]rowcodec.DefCausInfo, 0, len(e.columnInfos))
   181  	)
   182  	for i := range e.columnInfos {
   183  		info := e.columnInfos[i]
   184  		ft := e.fieldTps[i]
   185  		col := rowcodec.DefCausInfo{
   186  			ID:         info.DeferredCausetId,
   187  			Ft:         ft,
   188  			IsPKHandle: info.PkHandle,
   189  		}
   190  		defcaus = append(defcaus, col)
   191  		if info.PkHandle {
   192  			pkDefCauss = append(pkDefCauss, info.DeferredCausetId)
   193  		}
   194  	}
   195  	if len(pkDefCauss) == 0 {
   196  		if e.primaryDefCauss != nil {
   197  			pkDefCauss = e.primaryDefCauss
   198  		} else {
   199  			pkDefCauss = []int64{0}
   200  		}
   201  	}
   202  	def := func(i int, chk *chunk.Chunk) error {
   203  		info := e.columnInfos[i]
   204  		if info.PkHandle || len(info.DefaultVal) == 0 {
   205  			chk.AppendNull(i)
   206  			return nil
   207  		}
   208  		causetDecoder := codec.NewCausetDecoder(chk, e.sc.TimeZone)
   209  		_, err := causetDecoder.DecodeOne(info.DefaultVal, i, e.fieldTps[i])
   210  		if err != nil {
   211  			return err
   212  		}
   213  		return nil
   214  	}
   215  	return rowcodec.NewChunkCausetDecoder(defcaus, pkDefCauss, def, e.sc.TimeZone), nil
   216  }
   217  
   218  // decodeRelatedDeferredCausetVals decodes data to Causet slice according to the event information.
   219  func (e *evalContext) decodeRelatedDeferredCausetVals(relatedDefCausOffsets []int, value [][]byte, event []types.Causet) error {
   220  	var err error
   221  	for _, offset := range relatedDefCausOffsets {
   222  		event[offset], err = blockcodec.DecodeDeferredCausetValue(value[offset], e.fieldTps[offset], e.sc.TimeZone)
   223  		if err != nil {
   224  			return errors.Trace(err)
   225  		}
   226  	}
   227  	return nil
   228  }
   229  
   230  // flagsToStatementContext creates a StatementContext from a `fidelpb.SelectRequest.Flags`.
   231  func flagsToStatementContext(flags uint64) *stmtctx.StatementContext {
   232  	sc := new(stmtctx.StatementContext)
   233  	sc.IgnoreTruncate = (flags & perceptron.FlagIgnoreTruncate) > 0
   234  	sc.TruncateAsWarning = (flags & perceptron.FlagTruncateAsWarning) > 0
   235  	sc.InInsertStmt = (flags & perceptron.FlagInInsertStmt) > 0
   236  	sc.InSelectStmt = (flags & perceptron.FlagInSelectStmt) > 0
   237  	sc.InDeleteStmt = (flags & perceptron.FlagInUFIDelateOrDeleteStmt) > 0
   238  	sc.OverflowAsWarning = (flags & perceptron.FlagOverflowAsWarning) > 0
   239  	sc.IgnoreZeroInDate = (flags & perceptron.FlagIgnoreZeroInDate) > 0
   240  	sc.DividedByZeroAsWarning = (flags & perceptron.FlagDividedByZeroAsWarning) > 0
   241  	return sc
   242  }
   243  
   244  // ErrLocked is returned when trying to Read/Write on a locked key. Client should
   245  // backoff or cleanup the dagger then retry.
   246  type ErrLocked struct {
   247  	Key      []byte
   248  	Primary  []byte
   249  	StartTS  uint64
   250  	TTL      uint64
   251  	LockType uint8
   252  }
   253  
   254  // BuildLockErr generates ErrKeyLocked objects
   255  func BuildLockErr(key []byte, primaryKey []byte, startTS uint64, TTL uint64, lockType uint8) *ErrLocked {
   256  	errLocked := &ErrLocked{
   257  		Key:      key,
   258  		Primary:  primaryKey,
   259  		StartTS:  startTS,
   260  		TTL:      TTL,
   261  		LockType: lockType,
   262  	}
   263  	return errLocked
   264  }
   265  
   266  // Error formats the dagger to a string.
   267  func (e *ErrLocked) Error() string {
   268  	return fmt.Sprintf("key is locked, key: %q, Type: %v, primary: %q, startTS: %v", e.Key, e.LockType, e.Primary, e.StartTS)
   269  }
   270  
   271  func buildResp(chunks []fidelpb.Chunk, counts []int64, posetPosetDagReq *fidelpb.PosetDagRequest, err error, warnings []stmtctx.ALLEGROSQLWarn, dur time.Duration) *interlock.Response {
   272  	resp := &interlock.Response{}
   273  	selResp := &fidelpb.SelectResponse{
   274  		Error:        toPBError(err),
   275  		Chunks:       chunks,
   276  		OutputCounts: counts,
   277  	}
   278  	if posetPosetDagReq.DefCauslectInterDircutionSummaries != nil && *posetPosetDagReq.DefCauslectInterDircutionSummaries {
   279  		execSummary := make([]*fidelpb.InterlockingDirectorateInterDircutionSummary, len(posetPosetDagReq.InterlockingDirectorates))
   280  		for i := range execSummary {
   281  			// TODO: Add real interlock execution summary information.
   282  			execSummary[i] = &fidelpb.InterlockingDirectorateInterDircutionSummary{}
   283  		}
   284  		selResp.InterDircutionSummaries = execSummary
   285  	}
   286  	if len(warnings) > 0 {
   287  		selResp.Warnings = make([]*fidelpb.Error, 0, len(warnings))
   288  		for i := range warnings {
   289  			selResp.Warnings = append(selResp.Warnings, toPBError(warnings[i].Err))
   290  		}
   291  	}
   292  	if locked, ok := errors.Cause(err).(*ErrLocked); ok {
   293  		resp.Locked = &ekvrpcpb.LockInfo{
   294  			Key:         locked.Key,
   295  			PrimaryLock: locked.Primary,
   296  			LockVersion: locked.StartTS,
   297  			LockTtl:     locked.TTL,
   298  		}
   299  	}
   300  	resp.InterDircDetails = &ekvrpcpb.InterDircDetails{
   301  		HandleTime: &ekvrpcpb.HandleTime{ProcessMs: int64(dur / time.Millisecond)},
   302  	}
   303  	data, err := proto.Marshal(selResp)
   304  	if err != nil {
   305  		resp.OtherError = err.Error()
   306  		return resp
   307  	}
   308  	resp.Data = data
   309  	return resp
   310  }
   311  
   312  func toPBError(err error) *fidelpb.Error {
   313  	if err == nil {
   314  		return nil
   315  	}
   316  	perr := new(fidelpb.Error)
   317  	e := errors.Cause(err)
   318  	switch y := e.(type) {
   319  	case *terror.Error:
   320  		tmp := terror.ToALLEGROSQLError(y)
   321  		perr.Code = int32(tmp.Code)
   322  		perr.Msg = tmp.Message
   323  	case *allegrosql.ALLEGROSQLError:
   324  		perr.Code = int32(y.Code)
   325  		perr.Msg = y.Message
   326  	default:
   327  		perr.Code = int32(1)
   328  		perr.Msg = err.Error()
   329  	}
   330  	return perr
   331  }
   332  
   333  // extractKVRanges extracts ekv.KeyRanges slice from a SelectRequest.
   334  func extractKVRanges(startKey, endKey []byte, keyRanges []*interlock.KeyRange, descScan bool) (ekvRanges []ekv.KeyRange, err error) {
   335  	ekvRanges = make([]ekv.KeyRange, 0, len(keyRanges))
   336  	for _, kran := range keyRanges {
   337  		if bytes.Compare(kran.GetStart(), kran.GetEnd()) >= 0 {
   338  			err = errors.Errorf("invalid range, start should be smaller than end: %v %v", kran.GetStart(), kran.GetEnd())
   339  			return
   340  		}
   341  
   342  		upperKey := kran.GetEnd()
   343  		if bytes.Compare(upperKey, startKey) <= 0 {
   344  			continue
   345  		}
   346  		lowerKey := kran.GetStart()
   347  		if len(endKey) != 0 && bytes.Compare(lowerKey, endKey) >= 0 {
   348  			break
   349  		}
   350  		r := ekv.KeyRange{
   351  			StartKey: ekv.Key(maxStartKey(lowerKey, startKey)),
   352  			EndKey:   ekv.Key(minEndKey(upperKey, endKey)),
   353  		}
   354  		ekvRanges = append(ekvRanges, r)
   355  	}
   356  	if descScan {
   357  		reverseKVRanges(ekvRanges)
   358  	}
   359  	return
   360  }
   361  
   362  func reverseKVRanges(ekvRanges []ekv.KeyRange) {
   363  	for i := 0; i < len(ekvRanges)/2; i++ {
   364  		j := len(ekvRanges) - i - 1
   365  		ekvRanges[i], ekvRanges[j] = ekvRanges[j], ekvRanges[i]
   366  	}
   367  }
   368  
   369  func maxStartKey(rangeStartKey ekv.Key, regionStartKey []byte) []byte {
   370  	if bytes.Compare([]byte(rangeStartKey), regionStartKey) > 0 {
   371  		return []byte(rangeStartKey)
   372  	}
   373  	return regionStartKey
   374  }
   375  
   376  func minEndKey(rangeEndKey ekv.Key, regionEndKey []byte) []byte {
   377  	if len(regionEndKey) == 0 || bytes.Compare([]byte(rangeEndKey), regionEndKey) < 0 {
   378  		return []byte(rangeEndKey)
   379  	}
   380  	return regionEndKey
   381  }
   382  
   383  const rowsPerChunk = 64
   384  
   385  func appendRow(chunks []fidelpb.Chunk, data []byte, rowCnt int) []fidelpb.Chunk {
   386  	if rowCnt%rowsPerChunk == 0 {
   387  		chunks = append(chunks, fidelpb.Chunk{})
   388  	}
   389  	cur := &chunks[len(chunks)-1]
   390  	cur.RowsData = append(cur.RowsData, data...)
   391  	return chunks
   392  }
   393  
   394  // fieldTypeFromPBDeferredCauset creates a types.FieldType from fidelpb.DeferredCausetInfo.
   395  func fieldTypeFromPBDeferredCauset(col *fidelpb.DeferredCausetInfo) *types.FieldType {
   396  	return &types.FieldType{
   397  		Tp:          byte(col.GetTp()),
   398  		Flag:        uint(col.Flag),
   399  		Flen:        int(col.GetDeferredCausetLen()),
   400  		Decimal:     int(col.GetDecimal()),
   401  		Elems:       col.Elems,
   402  		DefCauslate: allegrosql.DefCauslations[uint8(collate.RestoreDefCauslationIDIfNeeded(col.GetDefCauslation()))],
   403  	}
   404  }
   405  
   406  // handleCopChecksumRequest handles interlock check sum request.
   407  func handleCopChecksumRequest(dbReader *dbreader.DBReader, req *interlock.Request) *interlock.Response {
   408  	resp := &fidelpb.ChecksumResponse{
   409  		Checksum:   1,
   410  		TotalEkvs:  1,
   411  		TotalBytes: 1,
   412  	}
   413  	data, err := resp.Marshal()
   414  	if err != nil {
   415  		return &interlock.Response{OtherError: fmt.Sprintf("marshal checksum response error: %v", err)}
   416  	}
   417  	return &interlock.Response{Data: data}
   418  }