github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/meta/internal/etcdkv/mock.go (about)

     1  // Copyright 2022 PingCAP, 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 etcdkv
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  	"net/url"
    20  	"os"
    21  	"strings"
    22  	"time"
    23  
    24  	"github.com/phayes/freeport"
    25  	"github.com/pingcap/failpoint"
    26  	"github.com/pingcap/log"
    27  	"github.com/pingcap/tiflow/pkg/errors"
    28  	"github.com/pingcap/tiflow/pkg/retry"
    29  	"go.etcd.io/etcd/server/v3/embed"
    30  	"go.uber.org/zap"
    31  )
    32  
    33  func allocTempURL() (string, error) {
    34  	port, err := freeport.GetFreePort()
    35  	if err != nil {
    36  		return "", err
    37  	}
    38  	return fmt.Sprintf("http://127.0.0.1:%d", port), nil
    39  }
    40  
    41  // RetryMockBackendEtcd retry to create backend DB if meet 'address already in use' error
    42  // for at most 3 times.
    43  func RetryMockBackendEtcd() (s *embed.Etcd, addr string, err error) {
    44  	err = retry.Do(context.TODO(), func() error {
    45  		s, addr, err = mockBackendEtcd()
    46  		if err != nil {
    47  			return err
    48  		}
    49  		return nil
    50  	},
    51  		retry.WithBackoffBaseDelay(1000 /* 1000 ms */),
    52  		retry.WithBackoffMaxDelay(3000 /* 3 seconds */),
    53  		retry.WithMaxTries(3 /* fail after 10 seconds*/),
    54  		retry.WithIsRetryableErr(func(err error) bool {
    55  			if strings.Contains(err.Error(), "address already in use") {
    56  				log.Info("address already in use, retry again")
    57  				return true
    58  			}
    59  			return false
    60  		}),
    61  	)
    62  
    63  	return
    64  }
    65  
    66  // mockBackendEtcd mock the etcd using embedded etcd as backend storage
    67  func mockBackendEtcd() (*embed.Etcd, string, error) {
    68  	failpoint.Inject("MockEtcdAddressAlreadyUse", func() {
    69  		failpoint.Return(nil, "", errors.New("address already in use"))
    70  	})
    71  
    72  	cfg := embed.NewConfig()
    73  	tmpDir := "embedded-etcd"
    74  	dir, err := os.MkdirTemp("", tmpDir)
    75  	if err != nil {
    76  		return nil, "", err
    77  	}
    78  	cfg.Dir = dir
    79  	peers, err := allocTempURL()
    80  	if err != nil {
    81  		return nil, "", err
    82  	}
    83  	log.Info("Allocate server peer port", zap.String("peers", peers))
    84  	u, err := url.Parse(peers)
    85  	if err != nil {
    86  		return nil, "", err
    87  	}
    88  	cfg.ListenPeerUrls = []url.URL{*u}
    89  	advertises, err := allocTempURL()
    90  	if err != nil {
    91  		return nil, "", err
    92  	}
    93  	log.Info("Allocate server advertises port", zap.String("advertises", advertises))
    94  	u, err = url.Parse(advertises)
    95  	if err != nil {
    96  		return nil, "", err
    97  	}
    98  	cfg.ListenClientUrls = []url.URL{*u}
    99  	svr, err := embed.StartEtcd(cfg)
   100  	if err != nil {
   101  		return nil, "", err
   102  	}
   103  	select {
   104  	case <-svr.Server.ReadyNotify():
   105  		log.Info("Server is ready!")
   106  	case <-time.After(60 * time.Second):
   107  		svr.Server.Stop() // trigger a shutdown
   108  		svr.Close()
   109  		return nil, "", errors.New("embedded etcd start fail")
   110  	}
   111  
   112  	return svr, advertises, nil
   113  }
   114  
   115  // CloseEmbededEtcd close the input embedded etcd server
   116  func CloseEmbededEtcd(svr *embed.Etcd) {
   117  	if svr != nil {
   118  		svr.Server.Stop()
   119  		svr.Close()
   120  	}
   121  }