add a new option: WithH24Validator

zhangbiao2009 · Jun 19, 2018 · 6f604d3 · 6f604d3
1 parent 6f26fa8
commit 6f604d3
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 15 deletions.
diff --git a/callback/wuid.go b/callback/wuid.go
@@ -70,7 +70,7 @@ func (this *WUID) LoadH24WithCallback(cb func() (uint64, error)) error {
 	}
 
 	this.w.Reset(h24 << 40)
-	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d", h24))
+	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d. tag: %s", h24, this.w.Tag))
 
 	this.w.Lock()
 	defer this.w.Unlock()
@@ -98,3 +98,8 @@ type Option internal.Option
 func WithSection(section uint8) Option {
 	return Option(internal.WithSection(section))
 }
+
+// WithH24Validator sets your own h24 validator
+func WithH24Validator(cb func(h24 uint64) error) Option {
+	return Option(internal.WithH24Validator(cb))
+}
diff --git a/internal/wuid.go b/internal/wuid.go
@@ -20,11 +20,12 @@ const (
 // WUID is for internal use only.
 type WUID struct {
 	sync.Mutex
-	Section uint8
-	N       uint64
-	Tag     string
-	Logger  Logger
-	Renew   func() error
+	Section      uint8
+	N            uint64
+	Tag          string
+	Logger       Logger
+	Renew        func() error
+	H24Validator func(h24 uint64) error
 }
 
 // NewWUID is for internal use only.
@@ -45,7 +46,7 @@ func NewWUID(tag string, logger Logger, opts ...Option) *WUID {
 func (this *WUID) Next() uint64 {
 	x := atomic.AddUint64(&this.N, 1)
 	if x&0xFFFFFFFFFF >= PanicValue {
-		panic(errors.New("<wuid> the low 40 bits are about to run out"))
+		panic("<wuid> the low 40 bits are about to run out")
 	}
 	if x&0xFFFFFFFFFF >= CriticalValue && x&RenewInterval == 0 {
 		go func() {
@@ -57,7 +58,7 @@ func (this *WUID) Next() uint64 {
 
 			err := this.RenewNow()
 			if err != nil {
-				this.Logger.Warn(fmt.Sprintf("<wuid> renew failed. tag: %s, reason: %s", this.Tag, err.Error()))
+				this.Logger.Warn(fmt.Sprintf("<wuid> renew failed. reason: %s", err.Error()))
 			} else {
 				this.Logger.Info(fmt.Sprintf("<wuid> renew succeeded. tag: %s", this.Tag))
 			}
@@ -100,6 +101,12 @@ func (this *WUID) VerifyH24(h24 uint64) error {
 		}
 	}
 
+	if this.H24Validator != nil {
+		if err := this.H24Validator(h24); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -131,3 +138,10 @@ func WithSection(section uint8) Option {
 		w.Section = section
 	}
 }
+
+// WithH24Validator is for internal use only.
+func WithH24Validator(cb func(h24 uint64) error) Option {
+	return func(w *WUID) {
+		w.H24Validator = cb
+	}
+}
diff --git a/internal/wuid_test.go b/internal/wuid_test.go
@@ -30,7 +30,7 @@ func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }
 func (p uint64Slice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
 
 func TestWUID_Next_Concurrent(t *testing.T) {
-	const total = 100
+	const total = 1000
 	g := NewWUID("default", nil)
 	var m sync.Mutex
 	var a = make(uint64Slice, 0, total)
@@ -212,3 +212,15 @@ func TestWithSection_Reset(t *testing.T) {
 		}
 	}
 }
+
+func TestWithRenewCallback(t *testing.T) {
+	g := NewWUID("default", nil, WithH24Validator(func(h24 uint64) error {
+		if h24 >= 10 {
+			return errors.New("bomb")
+		}
+		return nil
+	}))
+	if err := g.VerifyH24(10); err.Error() != "bomb" {
+		t.Fatal("the H24Validator was not called")
+	}
+}
diff --git a/mongo/wuid.go b/mongo/wuid.go
@@ -96,7 +96,7 @@ func (this *WUID) LoadH24FromMongoWithTimeout(addr, user, pass, dbName, coll, do
 	}
 
 	this.w.Reset(h24 << 40)
-	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d", h24))
+	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d. tag: %s", h24, this.w.Tag))
 
 	this.w.Lock()
 	defer this.w.Unlock()
@@ -124,3 +124,8 @@ type Option internal.Option
 func WithSection(section uint8) Option {
 	return Option(internal.WithSection(section))
 }
+
+// WithH24Validator sets your own h24 validator
+func WithH24Validator(cb func(h24 uint64) error) Option {
+	return Option(internal.WithH24Validator(cb))
+}
diff --git a/mysql/wuid.go b/mysql/wuid.go
@@ -87,7 +87,7 @@ func (this *WUID) LoadH24FromMysql(addr, user, pass, dbName, table string) error
 	}
 
 	this.w.Reset(h24 << 40)
-	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d", h24))
+	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d. tag: %s", h24, this.w.Tag))
 
 	this.w.Lock()
 	defer this.w.Unlock()
@@ -115,3 +115,8 @@ type Option internal.Option
 func WithSection(section uint8) Option {
 	return Option(internal.WithSection(section))
 }
+
+// WithH24Validator sets your own h24 validator
+func WithH24Validator(cb func(h24 uint64) error) Option {
+	return Option(internal.WithH24Validator(cb))
+}
diff --git a/pgsql/wuid.go b/pgsql/wuid.go
@@ -111,7 +111,7 @@ func (this *WUID) loadH24FromPg(dsn, table string) error {
 
 	db, err := sql.Open("postgres", dsn)
 	if err != nil {
-		return fmt.Errorf("db connection error: %s , with connection: %s", err, dsn)
+		return fmt.Errorf("db connection error: %s , with connection: %s, tag: %s", err, dsn, this.w.Tag)
 	}
 	defer db.Close()
 
@@ -127,7 +127,7 @@ func (this *WUID) loadH24FromPg(dsn, table string) error {
 	}
 
 	this.w.Reset(h24 << 40)
-	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d", h24))
+	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d. tag: %s", h24, this.w.Tag))
 
 	this.w.Lock()
 	defer this.w.Unlock()
@@ -155,3 +155,8 @@ type Option internal.Option
 func WithSection(section uint8) Option {
 	return Option(internal.WithSection(section))
 }
+
+// WithH24Validator sets your own h24 validator
+func WithH24Validator(cb func(h24 uint64) error) Option {
+	return Option(internal.WithH24Validator(cb))
+}
diff --git a/redis/wuid.go b/redis/wuid.go
@@ -69,7 +69,7 @@ func (this *WUID) LoadH24FromRedis(addr, pass, key string) error {
 	}
 
 	this.w.Reset(h24 << 40)
-	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d", h24))
+	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d. tag: %s", h24, this.w.Tag))
 
 	this.w.Lock()
 	defer this.w.Unlock()
@@ -110,7 +110,7 @@ func (this *WUID) LoadH24FromRedisCluster(addrs []string, pass, key string) erro
 	}
 
 	this.w.Reset(h24 << 40)
-	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d", h24))
+	this.w.Logger.Info(fmt.Sprintf("<wuid> new h24: %d. tag: %s", h24, this.w.Tag))
 
 	this.w.Lock()
 	defer this.w.Unlock()
@@ -138,3 +138,8 @@ type Option internal.Option
 func WithSection(section uint8) Option {
 	return Option(internal.WithSection(section))
 }
+
+// WithH24Validator sets your own h24 validator
+func WithH24Validator(cb func(h24 uint64) error) Option {
+	return Option(internal.WithH24Validator(cb))
+}