diff --git a/pkg/utils/tempurl/tempurl.go b/pkg/utils/tempurl/tempurl.go index 421513ff001..1ca3a6561af 100644 --- a/pkg/utils/tempurl/tempurl.go +++ b/pkg/utils/tempurl/tempurl.go @@ -16,7 +16,9 @@ package tempurl import ( "fmt" + "io" "net" + "net/http" "time" "github.com/pingcap/log" @@ -29,6 +31,9 @@ var ( testAddrMap = make(map[string]struct{}) ) +// reference: /pd/tools/pd-ut/alloc/server.go +const utAllocURL = "http://0.0.0.0:20180/alloc" + // Alloc allocates a local URL for testing. func Alloc() string { for i := 0; i < 10; i++ { @@ -42,6 +47,9 @@ func Alloc() string { } func tryAllocTestURL() string { + if url := getFromUT(); url != "" { + return url + } l, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { log.Fatal("listen failed", errs.ZapError(err)) @@ -63,3 +71,18 @@ func tryAllocTestURL() string { testAddrMap[addr] = struct{}{} return addr } + +func getFromUT() string { + resp, err := http.Get(utAllocURL) + if err != nil || resp.StatusCode != http.StatusOK { + return "" + } + + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return "" + } + url := string(body) + return url +} diff --git a/tests/server/member/member_test.go b/tests/server/member/member_test.go index c581eb39390..b13705bc534 100644 --- a/tests/server/member/member_test.go +++ b/tests/server/member/member_test.go @@ -289,7 +289,7 @@ func TestMoveLeader(t *testing.T) { re := require.New(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() - cluster, err := tests.NewTestCluster(ctx, 5) + cluster, err := tests.NewTestCluster(ctx, 3) defer cluster.Destroy() re.NoError(err) @@ -298,7 +298,7 @@ func TestMoveLeader(t *testing.T) { cluster.WaitLeader() var wg sync.WaitGroup - wg.Add(5) + wg.Add(3) for _, s := range cluster.GetServers() { go func(s *tests.TestServer) { defer wg.Done() diff --git a/tools/go.mod b/tools/go.mod index 2febbe1ad68..15b70979f7c 100644 --- a/tools/go.mod +++ b/tools/go.mod @@ -9,6 +9,7 @@ replace ( require ( github.com/BurntSushi/toml v0.3.1 + github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5 github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e github.com/coreos/go-semver v0.3.1 github.com/docker/go-units v0.4.0 @@ -64,7 +65,6 @@ require ( github.com/bitly/go-simplejson v0.5.0 // indirect github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch // indirect github.com/bytedance/sonic v1.9.1 // indirect - github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5 // indirect github.com/cenkalti/backoff/v4 v4.0.2 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect diff --git a/tools/pd-ut/alloc/check_env_dummy.go b/tools/pd-ut/alloc/check_env_dummy.go new file mode 100644 index 00000000000..b9b8eb4827a --- /dev/null +++ b/tools/pd-ut/alloc/check_env_dummy.go @@ -0,0 +1,21 @@ +// Copyright 2024 TiKV Project Authors. +// +// 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. +//go:build !linux +// +build !linux + +package alloc + +func environmentCheck(_ string) bool { + return true +} diff --git a/tools/pd-ut/alloc/check_env_linux.go b/tools/pd-ut/alloc/check_env_linux.go new file mode 100644 index 00000000000..1a51f8075cf --- /dev/null +++ b/tools/pd-ut/alloc/check_env_linux.go @@ -0,0 +1,42 @@ +// Copyright 2024 TiKV Project Authors. +// +// 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. +//go:build linux +// +build linux + +package alloc + +import ( + "github.com/cakturk/go-netstat/netstat" + "github.com/pingcap/log" + "go.uber.org/zap" +) + +func environmentCheck(addr string) bool { + valid, err := checkAddr(addr[len("http://"):]) + if err != nil { + log.Error("check port status failed", zap.Error(err)) + return false + } + return valid +} + +func checkAddr(addr string) (bool, error) { + tabs, err := netstat.TCPSocks(func(s *netstat.SockTabEntry) bool { + return s.RemoteAddr.String() == addr || s.LocalAddr.String() == addr + }) + if err != nil { + return false, err + } + return len(tabs) < 1, nil +} diff --git a/tools/pd-ut/alloc/server.go b/tools/pd-ut/alloc/server.go new file mode 100644 index 00000000000..42c6bbe3795 --- /dev/null +++ b/tools/pd-ut/alloc/server.go @@ -0,0 +1,48 @@ +// Copyright 2024 TiKV Project Authors. +// +// 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 alloc + +import ( + "errors" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/pingcap/log" + flag "github.com/spf13/pflag" + "go.uber.org/zap" +) + +var statusAddress = flag.String("status-addr", "0.0.0.0:20180", "status address") + +func RunHTTPServer() *http.Server { + gin.SetMode(gin.ReleaseMode) + engine := gin.New() + engine.Use(gin.Recovery()) + + engine.GET("alloc", func(c *gin.Context) { + addr := Alloc() + c.String(http.StatusOK, addr) + }) + + srv := &http.Server{Addr: *statusAddress, Handler: engine.Handler(), ReadHeaderTimeout: 3 * time.Second} + go func() { + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.Fatal("server listen error", zap.Error(err)) + } + }() + + return srv +} diff --git a/tools/pd-ut/alloc/tempurl.go b/tools/pd-ut/alloc/tempurl.go new file mode 100644 index 00000000000..6be69dfe056 --- /dev/null +++ b/tools/pd-ut/alloc/tempurl.go @@ -0,0 +1,65 @@ +// Copyright 2024 TiKV Project Authors. +// +// 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 alloc + +import ( + "fmt" + "net" + "sync" + "time" + + "github.com/pingcap/log" + "github.com/tikv/pd/pkg/errs" +) + +var ( + testAddrMutex sync.Mutex + testAddrMap = make(map[string]struct{}) +) + +// Alloc allocates a local URL for testing. +func Alloc() string { + for i := 0; i < 50; i++ { + if u := tryAllocTestURL(); u != "" { + return u + } + time.Sleep(200 * time.Millisecond) + } + log.Fatal("failed to alloc test URL") + return "" +} + +func tryAllocTestURL() string { + l, err := net.Listen("tcp", "127.0.0.1:") + if err != nil { + return "" + } + addr := fmt.Sprintf("http://%s", l.Addr()) + err = l.Close() + if err != nil { + log.Fatal("close failed", errs.ZapError(err)) + } + + testAddrMutex.Lock() + defer testAddrMutex.Unlock() + if _, ok := testAddrMap[addr]; ok { + return "" + } + if !environmentCheck(addr) { + return "" + } + testAddrMap[addr] = struct{}{} + return addr +} diff --git a/tools/pd-ut/ut.go b/tools/pd-ut/ut.go index 9419363c152..7781ab4ee3b 100644 --- a/tools/pd-ut/ut.go +++ b/tools/pd-ut/ut.go @@ -16,6 +16,7 @@ package main import ( "bytes" + "context" "encoding/xml" "errors" "fmt" @@ -32,6 +33,9 @@ import ( "sync" "time" + "github.com/tikv/pd/tools/pd-ut/alloc" + "go.uber.org/zap" + // Set the correct value when it runs inside docker. _ "go.uber.org/automaxprocs" ) @@ -128,6 +132,13 @@ func main() { fmt.Println("os.Getwd() error", err) } + srv := alloc.RunHTTPServer() + defer func() { + if err := srv.Shutdown(context.Background()); err != nil { + log.Fatal("server shutdown error", zap.Error(err)) + } + }() + var isSucceed bool // run all tests if len(os.Args) == 1 { @@ -684,7 +695,7 @@ func buildTestBinaryMulti(pkgs []string) error { } // go test --exec=xprog --tags=tso_function_test,deadlock -vet=off --count=0 $(pkgs) - // workPath just like `/data/nvme0n1/husharp/proj/pd/tests/integrations` + // workPath just like `/pd/tests/integrations` xprogPath := path.Join(workDir, "bin/xprog") if strings.Contains(workDir, integrationsTestPath) { xprogPath = path.Join(workDir[:strings.LastIndex(workDir, integrationsTestPath)], "bin/xprog")