tools/ut: add alloc server for pd-ut (#8203)

ref #7969

Signed-off-by: husharp <[email protected]>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

pkg/utils/tempurl/tempurl.go

@@ -16,7 +16,10 @@ package tempurl
 
 import (
 	"fmt"
+	"io"
 	"net"
+	"net/http"
+	"os"
 	"time"
 
 	"github.com/pingcap/log"
@@ -29,6 +32,9 @@ var (
 	testAddrMap   = make(map[string]struct{})
 )
 
+// reference: /pd/tools/pd-ut/alloc/server.go
+const AllocURLFromUT = "allocURLFromUT"
+
 // Alloc allocates a local URL for testing.
 func Alloc() string {
 	for i := 0; i < 10; i++ {
@@ -42,6 +48,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 +72,26 @@ func tryAllocTestURL() string {
 	testAddrMap[addr] = struct{}{}
 	return addr
 }
+
+func getFromUT() string {
+	addr := os.Getenv(AllocURLFromUT)
+	if addr == "" {
+		return ""
+	}
+
+	req, err := http.NewRequest(http.MethodGet, addr, nil)
+	if err != nil {
+		return ""
+	}
+	resp, err := http.DefaultClient.Do(req)
+	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
+}

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

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
+}

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
+}

tools/pd-ut/alloc/server.go

@@ -0,0 +1,56 @@
+// 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"
+	"flag"
+	"fmt"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pingcap/log"
+	"github.com/tikv/pd/pkg/utils/tempurl"
+	"go.uber.org/zap"
+)
+
+var statusAddress = flag.String("status-addr", "0.0.0.0:20180", "status address")
+
+func RunHTTPServer() *http.Server {
+	err := os.Setenv(tempurl.AllocURLFromUT, fmt.Sprintf("http://%s/alloc", *statusAddress))
+	if err != nil {
+		fmt.Println(err)
+	}
+
+	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
+}

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
+}

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")