github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/pdutil/utils.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package pdutil
     4  
     5  import (
     6  	"bytes"
     7  	"context"
     8  	"crypto/tls"
     9  	"encoding/hex"
    10  	"encoding/json"
    11  	"fmt"
    12  	"net/http"
    13  	"strings"
    14  
    15  	"github.com/pingcap/errors"
    16  	"github.com/pingcap/tidb/tablecodec"
    17  	"github.com/tikv/pd/pkg/codec"
    18  	"github.com/tikv/pd/server/schedule/placement"
    19  
    20  	berrors "github.com/pingcap/br/pkg/errors"
    21  	"github.com/pingcap/br/pkg/httputil"
    22  )
    23  
    24  // UndoFunc is a 'undo' operation of some undoable command.
    25  // (e.g. RemoveSchedulers).
    26  type UndoFunc func(context.Context) error
    27  
    28  // Nop is the 'zero value' of undo func.
    29  var Nop UndoFunc = func(context.Context) error { return nil }
    30  
    31  const (
    32  	resetTSURL       = "/pd/api/v1/admin/reset-ts"
    33  	placementRuleURL = "/pd/api/v1/config/rules"
    34  )
    35  
    36  // ResetTS resets the timestamp of PD to a bigger value.
    37  func ResetTS(ctx context.Context, pdAddr string, ts uint64, tlsConf *tls.Config) error {
    38  	payload, err := json.Marshal(struct {
    39  		TSO string `json:"tso,omitempty"`
    40  	}{TSO: fmt.Sprintf("%d", ts)})
    41  	if err != nil {
    42  		return errors.Trace(err)
    43  	}
    44  	cli := httputil.NewClient(tlsConf)
    45  	prefix := "http://"
    46  	if tlsConf != nil {
    47  		prefix = "https://"
    48  	}
    49  	reqURL := prefix + pdAddr + resetTSURL
    50  	req, err := http.NewRequestWithContext(ctx, "POST", reqURL, strings.NewReader(string(payload)))
    51  	if err != nil {
    52  		return errors.Trace(err)
    53  	}
    54  	req.Header.Set("Content-Type", "application/json")
    55  	resp, err := cli.Do(req)
    56  	if err != nil {
    57  		return errors.Trace(err)
    58  	}
    59  	defer resp.Body.Close()
    60  	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusForbidden {
    61  		buf := new(bytes.Buffer)
    62  		_, _ = buf.ReadFrom(resp.Body)
    63  		return errors.Annotatef(berrors.ErrPDInvalidResponse, "pd resets TS failed: req=%v, resp=%v, err=%v", string(payload), buf.String(), err)
    64  	}
    65  	return nil
    66  }
    67  
    68  // GetPlacementRules return the current placement rules.
    69  func GetPlacementRules(ctx context.Context, pdAddr string, tlsConf *tls.Config) ([]placement.Rule, error) {
    70  	cli := httputil.NewClient(tlsConf)
    71  	prefix := "http://"
    72  	if tlsConf != nil {
    73  		prefix = "https://"
    74  	}
    75  	reqURL := prefix + pdAddr + placementRuleURL
    76  	req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
    77  	if err != nil {
    78  		return nil, errors.Trace(err)
    79  	}
    80  	resp, err := cli.Do(req)
    81  	if err != nil {
    82  		return nil, errors.Trace(err)
    83  	}
    84  	defer resp.Body.Close()
    85  	buf := new(bytes.Buffer)
    86  	_, err = buf.ReadFrom(resp.Body)
    87  	if err != nil {
    88  		return nil, errors.Trace(err)
    89  	}
    90  	if resp.StatusCode == http.StatusPreconditionFailed {
    91  		return []placement.Rule{}, nil
    92  	}
    93  	if resp.StatusCode != http.StatusOK {
    94  		return nil, errors.Annotatef(berrors.ErrPDInvalidResponse, "get placement rules failed: resp=%v, err=%v, code=%d", buf.String(), err, resp.StatusCode)
    95  	}
    96  	var rules []placement.Rule
    97  	err = json.Unmarshal(buf.Bytes(), &rules)
    98  	if err != nil {
    99  		return nil, errors.Trace(err)
   100  	}
   101  	return rules, nil
   102  }
   103  
   104  // SearchPlacementRule returns the placement rule matched to the table or nil.
   105  func SearchPlacementRule(tableID int64, placementRules []placement.Rule, role placement.PeerRoleType) *placement.Rule {
   106  	for _, rule := range placementRules {
   107  		key, err := hex.DecodeString(rule.StartKeyHex)
   108  		if err != nil {
   109  			continue
   110  		}
   111  		_, decoded, err := codec.DecodeBytes(key)
   112  		if err != nil {
   113  			continue
   114  		}
   115  		if rule.Role == role && tableID == tablecodec.DecodeTableID(decoded) {
   116  			return &rule
   117  		}
   118  	}
   119  	return nil
   120  }