add NextIds support

Author: Terry-Mao

CHANGELOG.md

@@ -1,4 +1,14 @@
-## Version 1.2 (Pennding)
+## Version 1.3 (Pennding)
+## Version 1.2 
+
+New Features:
+    - Add NextIds for get mutiple ids.
+    - Add "start"(start timestamp) config for calculating snowflake id.
+
+Bugfixes:
+
+ - 
+
 ## Version 1.1 (2014-12-15)
 
 New Features:

client/client.go

@@ -4,6 +4,7 @@ import (
 	log "code.google.com/p/log4go"
 	"encoding/json"
 	"errors"
+	myrpc "github.com/Terry-Mao/gosnowflake/rpc"
 	"github.com/samuel/go-zookeeper/zk"
 	"math/rand"
 	"net/rpc"
@@ -20,8 +21,9 @@ const (
 	rpcClientPingSleep  = 1 * time.Second // rpc client ping need sleep
 	rpcClientRetrySleep = 1 * time.Second // rpc client retry connect need sleep
 
-	RPCPing   = "SnowflakeRPC.Ping"
-	RPCNextId = "SnowflakeRPC.NextId"
+	RPCPing    = "SnowflakeRPC.Ping"
+	RPCNextId  = "SnowflakeRPC.NextId"
+	RPCNextIds = "SnowflakeRPC.NextIds"
 )
 
 var (
@@ -105,6 +107,18 @@ func (c *Client) Id() (id int64, err error) {
 	return
 }
 
+// Ids generate a snowflake id.
+func (c *Client) Ids(num int) (ids []int64, err error) {
+	client, err := c.client()
+	if err != nil {
+		return
+	}
+	if err = client.Call(RPCNextIds, &myrpc.NextIdsArgs{WorkerId: c.workerId, Num: num}, &ids); err != nil {
+		log.Error("rpc.Call(\"%s\", %d, &id) error(%v)", RPCNextId, c.workerId, err)
+	}
+	return
+}
+
 // closeRpc close rpc resource.
 func closeRpc(clients []*rpc.Client, stop chan bool) {
 	// rpc

client/client_test.go

@@ -59,7 +59,12 @@ func Test(t *testing.T) {
 		if err != nil {
 			t.Error(err)
 		}
+		ids, err := c.Ids(5)
+		if err != nil {
+			t.Error(err)
+		}
 		fmt.Printf("gosnwoflake id: %d\n", id)
+		fmt.Printf("gosnwoflake ids: %d\n", ids)
 	}
 	c.Close()
 	// check global cache map

client/test.conf

@@ -20,7 +20,7 @@
 
 [base]
 # rpc address
-rpc.addr 10.33.56.165:8080
+rpc.addr 10.33.60.26:8080
 
 # worker id
 worker 1

config.go

@@ -41,9 +41,11 @@ type Config struct {
 	PprofBind    []string      `goconf:"base:pprof.bind:,"`
 	DatacenterId int64         `goconf:"snowflake:datacenter"`
 	WorkerId     []int64       `goconf:"snowflake:worker"`
+	Start        string        `goconf:"snowflake:start"`
 	ZKAddr       []string      `goconf:"zookeeper:addr"`
 	ZKTimeout    time.Duration `goconf:"zookeeper:timeout:time`
 	ZKPath       string        `goconf:"zookeeper:path"`
+	Twepoch      int64
 }
 
 func init() {
@@ -52,6 +54,7 @@ func init() {
 
 // Init init the configuration file.
 func InitConfig() (err error) {
+	var twepoch time.Time
 	MyConf = &Config{
 		PidFile:      "/tmp/gosnowflake.pid",
 		Dir:          "/dev/null",
@@ -61,6 +64,7 @@ func InitConfig() (err error) {
 		ThriftBind:   []string{"localhost:8081"},
 		DatacenterId: 0,
 		WorkerId:     []int64{0},
+		Start:        "2010-11-04 09:42:54",
 		ZKAddr:       []string{"localhost:2181"},
 		ZKTimeout:    time.Second * 15,
 		ZKPath:       "/gosnowflake-servers",
@@ -71,5 +75,10 @@ func InitConfig() (err error) {
 	if err = goConf.Unmarshal(MyConf); err != nil {
 		return
 	}
+	if twepoch, err = time.Parse("2006-01-02 15:04:05", MyConf.Start); err != nil {
+		return
+	} else {
+		MyConf.Twepoch = twepoch.UnixNano() / int64(time.Millisecond)
+	}
 	return
 }

gosnowflake-example.conf

@@ -125,3 +125,11 @@ datacenter 0
 # worker 0
 # worker 0,1,2
 worker 0,1,2
+
+# start set the timestamp for calculate the snowflake id, current timestamp 
+# minus start timestamp.
+# default value is 2010-11-04 09:42:54
+# Examples:
+#
+# start 2010-11-04 09:42:54
+start 2010-11-04 09:42:54

id.go

@@ -43,12 +43,13 @@ type IdWorker struct {
 	sequence      int64
 	lastTimestamp int64
 	workerId      int64
+	twepoch       int64
 	datacenterId  int64
 	mutex         *sync.Mutex
 }
 
 // NewIdWorker new a snowflake id generator object.
-func NewIdWorker(workerId, datacenterId int64) (*IdWorker, error) {
+func NewIdWorker(workerId, datacenterId int64, twepoch int64) (*IdWorker, error) {
 	idWorker := &IdWorker{}
 	if workerId > maxWorkerId || workerId < 0 {
 		log.Error("worker Id can't be greater than %d or less than 0", maxWorkerId)
@@ -62,6 +63,7 @@ func NewIdWorker(workerId, datacenterId int64) (*IdWorker, error) {
 	idWorker.datacenterId = datacenterId
 	idWorker.lastTimestamp = -1
 	idWorker.sequence = 0
+	idWorker.twepoch = twepoch
 	idWorker.mutex = &sync.Mutex{}
 	log.Debug("worker starting. timestamp left shift %d, datacenter id bits %d, worker id bits %d, sequence bits %d, workerid %d", timestampLeftShift, datacenterIdBits, workerIdBits, sequenceBits, workerId)
 	return idWorker, nil
@@ -99,5 +101,5 @@ func (id *IdWorker) NextId() (int64, error) {
 		id.sequence = 0
 	}
 	id.lastTimestamp = timestamp
-	return ((timestamp - twepoch) << timestampLeftShift) | (id.datacenterId << datacenterIdShift) | (id.workerId << workerIdShift) | id.sequence, nil
+	return ((timestamp - id.twepoch) << timestampLeftShift) | (id.datacenterId << datacenterIdShift) | (id.workerId << workerIdShift) | id.sequence, nil
 }

id_test.go

@@ -22,7 +22,7 @@ import (
 )
 
 func TestID(t *testing.T) {
-	id, err := NewIdWorker(0, 0)
+	id, err := NewIdWorker(0, 0, twepoch)
 	if err != nil {
 		log.Error("NewIdWorker(0, 0) error(%v)", err)
 		t.FailNow()
@@ -36,7 +36,7 @@ func TestID(t *testing.T) {
 }
 
 func BenchmarkID(b *testing.B) {
-	id, err := NewIdWorker(0, 0)
+	id, err := NewIdWorker(0, 0, twepoch)
 	if err != nil {
 		log.Error("NewIdWorker(0, 0) error(%v)", err)
 		b.FailNow()

rpc.go

@@ -20,11 +20,16 @@ import (
 	log "code.google.com/p/log4go"
 	"errors"
 	"fmt"
+	myrpc "github.com/Terry-Mao/gosnowflake/rpc"
 	"net"
 	"net/rpc"
 	"time"
 )
 
+const (
+	maxNextIdsNum = 100
+)
+
 type SnowflakeRPC struct {
 	idWorkers []*IdWorker
 }
@@ -41,7 +46,7 @@ func InitRPC() error {
 			log.Error("init workerId: %d already exists", workerId)
 			return fmt.Errorf("init workerId: %d exists", workerId)
 		}
-		idWorker, err := NewIdWorker(workerId, MyConf.DatacenterId)
+		idWorker, err := NewIdWorker(workerId, MyConf.DatacenterId, MyConf.Twepoch)
 		if err != nil {
 			log.Error("NewIdWorker(%d, %d) error(%v)", MyConf.DatacenterId, workerId)
 			return err
@@ -98,6 +103,37 @@ func (s *SnowflakeRPC) NextId(workerId int64, id *int64) error {
 	return nil
 }
 
+// NextIds generate specified num ids.
+func (s *SnowflakeRPC) NextIds(args *myrpc.NextIdsArgs, ids *[]int64) error {
+	if args == nil {
+		return errors.New("args is nil")
+	}
+	if args.WorkerId > maxWorkerId || args.WorkerId < 0 {
+		log.Error("worker Id can't be greater than %d or less than 0", maxWorkerId)
+		return errors.New(fmt.Sprintf("worker Id: %d error", args.WorkerId))
+	}
+	if args.Num > maxNextIdsNum || args.Num < 0 {
+		log.Error("num can't be greater than %d or less than 0", maxNextIdsNum)
+		return errors.New(fmt.Sprintf("num: %d error", args.Num))
+	}
+	if worker := s.idWorkers[args.WorkerId]; worker == nil {
+		log.Warn("workerId: %d not register", args.WorkerId)
+		return fmt.Errorf("snowflake workerId: %d don't register in this service", args.WorkerId)
+	} else {
+		tids := make([]int64, args.Num)
+		for i := 0; i < args.Num; i++ {
+			if tid, err := worker.NextId(); err != nil {
+				log.Error("worker.NextId() error(%v)", err)
+				return err
+			} else {
+				tids[i] = tid
+			}
+		}
+		*ids = tids
+	}
+	return nil
+}
+
 // DatacenterId return the services's datacenterId.
 func (s *SnowflakeRPC) DatacenterId(ignore int, dataCenterId *int64) error {
 	*dataCenterId = MyConf.DatacenterId

rpc/args.go

@@ -0,0 +1,6 @@
+package rpc
+
+type NextIdsArgs struct {
+	WorkerId int64 // snowflake worker id
+	Num      int   // batch next id number
+}