fix(keyviz): find nearest id if not found

pkg/keyvisual/decorator/tidb.go

@@ -6,6 +6,7 @@
 	"context"
 	"encoding/hex"
 	"fmt"
+	"sort"
 	"sync"
 	"time"
 
@@ -24,6 +25,7 @@
 		EtcdClient:    etcdClient,
 		tidbClient:    tidbClient,
 		SchemaVersion: -1,
+		tableInOrder:  &tableInOrder{},
 	}
 
 	lc.Append(fx.Hook{
@@ -52,14 +54,103 @@
 	EtcdClient *clientv3.Client
 
 	TableMap      sync.Map
+	tableInOrder  *tableInOrder
 	tidbClient    *tidb.Client
 	SchemaVersion int64
 	TidbAddress   []string
 }
 
+type tableInOrder struct {
+	mu     sync.RWMutex
+	tables []*tableDetail
+}
+
+// BuildFromTableMap build ordered tables from a table map.
+func (inOrder *tableInOrder) buildFromTableMap(m *sync.Map) {
+	if m == nil {
+		return
+	}
+
+	tables := []*tableDetail{}
+	m.Range(func(key, value interface{}) bool {
+		t := value.(*tableDetail)
+		tables = append(tables, t)
+		return true
+	})
+
+	sort.Sort(&tableSorter{tables: tables})
+
+	inOrder.mu.Lock()
+	defer inOrder.mu.Unlock()
+	inOrder.tables = tables
+}
+
+// FindOne will find first table detail which id is between [fromId, toId).
+// Returns nil if not found any.
+func (inOrder *tableInOrder) findOne(fromID, toID int64) *tableDetail {
+	if fromID >= toID {
+		return nil
+	}
+
+	inOrder.mu.RLock()
+	defer inOrder.mu.RUnlock()
+
+	tLen := len(inOrder.tables)
+	if tLen == 0 {
+		return nil
+	}
+
+	pivot := tLen / 2
+	left := 0
+	right := tLen
+	for pivot > left {
+		prevID := inOrder.tables[pivot-1].ID
+		id := inOrder.tables[pivot].ID
+		// find approaching id near the fromId
+		// table_1 ------- table_3 ------- table_5
+		//           	^
+		//       search table_2
+		// approaching result: table_3
+		if prevID < fromID && id >= fromID {
+			break
+		}
+
+		if id < fromID {
+			left = pivot
+		} else {
+			right = pivot
+		}
+		pivot = (left + right) / 2
+	}
+
+	id := inOrder.tables[pivot].ID
+	if id < fromID || id > toID {
+		return nil
+	}
+
+	return inOrder.tables[pivot]
+}
+
+type tableSorter struct {
+	tables []*tableDetail
+}
+
+func (ts *tableSorter) Len() int {
+	return len(ts.tables)
+}
+
+func (ts *tableSorter) Swap(i, j int) {
+	ts.tables[i], ts.tables[j] = ts.tables[j], ts.tables[i]
+}
+
+func (ts *tableSorter) Less(i, j int) bool {
+	return ts.tables[i].ID < ts.tables[j].ID
+}
+
 type tidbLabeler struct {
-	TableMap *sync.Map
-	Buffer   model.KeyInfoBuffer
+	TableMap     *sync.Map
+	tableInOrder *tableInOrder
+	Buffer       model.KeyInfoBuffer
 }
 
 func (s *tidbLabelStrategy) ReloadConfig(cfg *config.KeyVisualConfig) {}
@@ -79,7 +170,8 @@
 
 func (s *tidbLabelStrategy) NewLabeler() Labeler {
 	return &tidbLabeler{
-		TableMap: &s.TableMap,
+		TableMap:     &s.TableMap,
+		tableInOrder: s.tableInOrder,
 	}
 }
 
@@ -105,8 +197,8 @@
 // Label will parse the ID information of the table and index.
 func (e *tidbLabeler) Label(keys []string) []LabelKey {
 	labelKeys := make([]LabelKey, len(keys))
-	for i, key := range keys {
-		labelKeys[i] = e.label(key)
+	for i := 1; i < len(keys); i++ {
+		labelKeys[i-1] = e.label(keys[i-1], keys[i])
 	}
 
 	if keys[0] == "" {
@@ -120,30 +212,40 @@
 	return labelKeys
 }
 
-func (e *tidbLabeler) label(key string) (label LabelKey) {
-	keyBytes := region.Bytes(key)
-	label.Key = hex.EncodeToString(keyBytes)
-	keyInfo, _ := e.Buffer.DecodeKey(keyBytes)
+func (e *tidbLabeler) label(lKey, rKey string) (label LabelKey) {
+	lKeyBytes := region.Bytes(lKey)
+	label.Key = hex.EncodeToString(lKeyBytes)
+	lKeyInfo, _ := e.Buffer.DecodeKey(lKeyBytes)
 
-	isMeta, tableID := keyInfo.MetaOrTable()
+	isMeta, lTableID := lKeyInfo.MetaOrTable()
 	if isMeta {
 		label.Labels = append(label.Labels, "meta")
 		return
 	}
 
 	var detail *tableDetail
-	if v, ok := e.TableMap.Load(tableID); ok {
+	if v, ok := e.TableMap.Load(lTableID); ok {
 		detail = v.(*tableDetail)
 		label.Labels = append(label.Labels, detail.DB, detail.Name)
 	} else {
-		label.Labels = append(label.Labels, fmt.Sprintf("table_%d", tableID))
+		rKeyBytes := region.Bytes(rKey)
+		rKeyInfo, _ := e.Buffer.DecodeKey(rKeyBytes)
+		_, rTableID := rKeyInfo.MetaOrTable()
+		detail := e.tableInOrder.findOne(lTableID, rTableID)
+
+		if detail != nil {
+			label.Labels = append(label.Labels, detail.DB, detail.Name)
+			// can't find the row/index info if the table info was came from a range
+			return
+		}
+		label.Labels = append(label.Labels, fmt.Sprintf("table_%d", lTableID))
 	}
 
-	if isCommonHandle, rowID := keyInfo.RowInfo(); isCommonHandle {
+	if isCommonHandle, rowID := lKeyInfo.RowInfo(); isCommonHandle {
 		label.Labels = append(label.Labels, "row")
 	} else if rowID != 0 {
 		label.Labels = append(label.Labels, fmt.Sprintf("row_%d", rowID))
-	} else if indexID := keyInfo.IndexInfo(); indexID != 0 {
+	} else if indexID := lKeyInfo.IndexInfo(); indexID != 0 {
 		if detail == nil {
 			label.Labels = append(label.Labels, fmt.Sprintf("index_%d", indexID))
 		} else if name, ok := detail.Indices[indexID]; ok {

pkg/keyvisual/decorator/tidb_requests.go

@@ -104,6 +104,8 @@
 		}
 	}
 
+	s.tableInOrder.buildFromTableMap(&s.TableMap)
+
 	// update schema version
 	if updateSuccess {
 		s.SchemaVersion = schemaVersion

pkg/keyvisual/decorator/tidb_test.go

@@ -3,9 +3,57 @@
 package decorator
 
 import (
+	"sync"
+	"testing"
+
 	. "github.com/pingcap/check"
+	"github.com/stretchr/testify/require"
 )
 
 var _ = Suite(&testTiDBSuite{})
 
 type testTiDBSuite struct{}
+
+func TestTableInOrderBuild(t *testing.T) {
+	tableMap := sync.Map{}
+	tableInOrder := &tableInOrder{}
+
+	tableMap.Store(8, &tableDetail{ID: 8})
+	tableMap.Store(2, &tableDetail{ID: 2})
+	tableMap.Store(4, &tableDetail{ID: 4})
+	tableMap.Store(1, &tableDetail{ID: 1})
+
+	tableInOrder.buildFromTableMap(&tableMap)
+	tableIds := make([]int64, 0, len(tableInOrder.tables))
+	for _, table := range tableInOrder.tables {
+		tableIds = append(tableIds, table.ID)
+	}
+
+	require.Equal(t, []int64{1, 2, 4, 8}, tableIds)
+}
+
+func TestTableInOrderFindOne(t *testing.T) {
+	table := &tableInOrder{
+		tables: []*tableDetail{{ID: 1}, {ID: 2}, {ID: 4}, {ID: 8}},
+	}
+	require.Equal(t, table.findOne(1, 2).ID, int64(1))
+	require.Equal(t, table.findOne(2, 3).ID, int64(2))
+	require.Equal(t, table.findOne(3, 5).ID, int64(4))
+	require.Equal(t, table.findOne(8, 18).ID, int64(8))
+	require.Equal(t, table.findOne(3, 4).ID, int64(4))
+	require.Nil(t, table.findOne(8, 0))
+	require.Nil(t, table.findOne(8, 8))
+	require.Nil(t, table.findOne(80, 81))
+
+	table0 := &tableInOrder{
+		tables: []*tableDetail{},
+	}
+	require.Nil(t, table0.findOne(1, 2))
+
+	table1 := &tableInOrder{
+		tables: []*tableDetail{{ID: 2}},
+	}
+	require.Equal(t, table1.findOne(1, 2).ID, int64(2))
+	require.Nil(t, table1.findOne(0, 1))
+	require.Nil(t, table1.findOne(3, 4))
+}

scripts/lint.sh

@@ -26,7 +26,7 @@ if [ "${NEED_DOWNLOAD}" = true ]; then
 fi
 
 echo "+ Run lints for Go source code"
-${LINT_BIN} run --fix
+${LINT_BIN} run --fix --timeout 3m
 
 echo "+ Clean up go mod"
 go mod tidy