72 lines
2.2 KiB
Go
72 lines
2.2 KiB
Go
// Copyright 2023 PingCAP, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package disttaskutil
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
|
|
"github.com/pingcap/tidb/pkg/domain/infosync"
|
|
"github.com/pingcap/tidb/pkg/domain/serverinfo"
|
|
)
|
|
|
|
// GenerateExecID used to generate IP:port as exec_id value
|
|
// This function is used by distributed task execution to generate serverID string to
|
|
// correlated one subtask to on TiDB node to be executed.
|
|
func GenerateExecID(info *serverinfo.ServerInfo) string {
|
|
return net.JoinHostPort(info.IP, fmt.Sprintf("%d", info.Port))
|
|
}
|
|
|
|
// MatchServerInfo will check if the schedulerID matched in all serverInfos.
|
|
func MatchServerInfo(serverInfos []*serverinfo.ServerInfo, schedulerID string) bool {
|
|
return FindServerInfo(serverInfos, schedulerID) >= 0
|
|
}
|
|
|
|
// FindServerInfo will find the schedulerID in all serverInfos.
|
|
func FindServerInfo(serverInfos []*serverinfo.ServerInfo, schedulerID string) int {
|
|
for i, serverInfo := range serverInfos {
|
|
serverID := GenerateExecID(serverInfo)
|
|
if serverID == schedulerID {
|
|
return i
|
|
}
|
|
}
|
|
return -1
|
|
}
|
|
|
|
// GenerateSubtaskExecID generates the subTask execID.
|
|
func GenerateSubtaskExecID(ctx context.Context, id string) string {
|
|
serverInfos, err := infosync.GetAllServerInfo(ctx)
|
|
if err != nil || len(serverInfos) == 0 {
|
|
return ""
|
|
}
|
|
if serverNode, ok := serverInfos[id]; ok {
|
|
return GenerateExecID(serverNode)
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// GenerateSubtaskExecID4Test generates the subTask execID, only used in unit tests.
|
|
func GenerateSubtaskExecID4Test(id string) string {
|
|
serverInfos := infosync.MockGlobalServerInfoManagerEntry.GetAllServerInfo()
|
|
if len(serverInfos) == 0 {
|
|
return ""
|
|
}
|
|
if serverNode, ok := serverInfos[id]; ok {
|
|
return GenerateExecID(serverNode)
|
|
}
|
|
return ""
|
|
}
|