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 }