Merge pull request #3 from wangce1998/feature/batch-id

批量获取分布式ID

main/id_app.go

@@ -7,6 +7,12 @@ import (
 	"strconv"
 )
 
+type Response struct {
+	Code    int         `json:"code"`
+	Message string      `json:"message"`
+	Data    interface{} `json:"data"`
+}
+
 func newIDApp(basePath string) http.Handler {
 
 	gin.SetMode(gin.ReleaseMode)
@@ -20,6 +26,39 @@ func newIDApp(basePath string) http.Handler {
 		c.String(200, strconv.FormatInt(id, 10))
 	})
 
+	router.GET("/batch", func(c *gin.Context) {
+		gen := c.Query("gen")
+		numStr := c.Query("num")
+
+		var res Response
+
+		num, err := strconv.Atoi(numStr)
+		if err != nil {
+			res.Code = 422
+			res.Message = err.Error()
+			c.JSON(http.StatusOK, res)
+
+			return
+		}
+		if num < 1 {
+			res.Code = 422
+			res.Message = "num最小值应为1"
+			c.JSON(http.StatusOK, res)
+
+			return
+		}
+
+		ids := xid.GetIDS(gen, num)
+
+		res.Code = 200
+		res.Message = "success"
+		res.Data = map[string]interface{}{
+			"ids": ids,
+		}
+
+		c.JSON(http.StatusOK, res)
+	})
+
 	router.GET("/gen/:gen", func(c *gin.Context) {
 		gen := c.Param("gen")
 		id := xid.MultiIdGenerator(gen).Next()

node.go

@@ -11,6 +11,9 @@ type NodeAllocation interface {
 	DestroyNode(timeoutCtx context.Context)
 }
 
+// MaxBatchNum 最大批量生成数
+const MaxBatchNum = 1000
+
 var mu = &sync.Mutex{}
 var curNodeId = -1
 var idGenerators = map[string]IDGen{}
@@ -67,3 +70,19 @@ func MultiIdGenerator(gen string) IDGen {
 	idGenerators[gen] = idGen
 	return idGen
 }
+
+// GetIDS 获取多个ID
+func GetIDS(gen string, num int) []int64 {
+	if num > MaxBatchNum {
+		num = MaxBatchNum
+	}
+	if num < 1 {
+		num = 1
+	}
+	idGen := MultiIdGenerator(gen)
+	var ids = make([]int64, num)
+	for i:=0;i<num; i++ {
+		ids[i] = idGen.Next()
+	}
+	return ids
+}

node_test.go

@@ -0,0 +1,41 @@
+package xid
+
+import (
+	"testing"
+)
+
+func TestGetIDS(t *testing.T) {
+	type args struct {
+		gen string
+		num int
+	}
+	tests := []struct {
+		name string
+		args args
+		want int
+	}{
+		{"-1", args{"", -1}, 1},
+		{"0", args{"", 0}, 1},
+		{"1", args{"", 1}, 1},
+		{"10", args{"", 10}, 10},
+		{"1000", args{"", 1000}, 1000},
+		{"10000", args{"", 10000}, 1000},
+	}
+
+	Config("snake", NewNodeAllocationSingle())
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := GetIDS(tt.args.gen, tt.args.num)
+			if len(got) != tt.want && got[0] > 0 {
+				t.Errorf("GetIDS() = %v, want %v", got, tt.want)
+			}
+
+			for i := 0; i < len(got)-1; i++ {
+				if got[i] >= got[i+1] {
+					t.Errorf("GetIDS() = \n%v\n%v", got[i], got[i+1])
+				}
+			}
+		})
+	}
+}