add feature control the max subscriber per key

Terry-Mao · Terry-Mao · commit 4ee7348f4437 · 2013-09-16T22:44:07.000+08:00
diff --git a/config.go b/config.go
@@ -16,22 +16,23 @@ func init() {
 }
 
 type Config struct {
-	Addr               string `json:"addr"`
-	Port               int    `json:"port"`
-	Pprof              int    `json:"pprof"`
-	PprofAddr          string `json:"pprof_addr"`
-	PprofPort          int    `json:"pprof_port"`
-	PubAddr            string `json:"pub_addr"`
-	PubPort            int    `json:"pub_port"`
-	LongpollingTimeout int    `json:"longpolling_timeout"`
-	MessageTimeout     int    `json:"message_timeout"`
-	Log                string `json:"log"`
-	RedisNetwork       string `json:"redis_network"`
-	RedisAddr          string `json:"redis_addr"`
-	RedisTimeout       int    `json:"redis_timeout"`
-	RedisPoolSize      int    `json:"redis_poolsize"`
-	RedisMQSize        int    `json:"redis_mqsize"`
-	MaxProcs           int    `json:"max_procs"`
+	Addr                string `json:"addr"`
+	Port                int    `json:"port"`
+	Pprof               int    `json:"pprof"`
+	PprofAddr           string `json:"pprof_addr"`
+	PprofPort           int    `json:"pprof_port"`
+	PubAddr             string `json:"pub_addr"`
+	PubPort             int    `json:"pub_port"`
+	LongpollingTimeout  int    `json:"longpolling_timeout"`
+	MessageTimeout      int    `json:"message_timeout"`
+	Log                 string `json:"log"`
+	RedisNetwork        string `json:"redis_network"`
+	RedisAddr           string `json:"redis_addr"`
+	RedisTimeout        int    `json:"redis_timeout"`
+	RedisPoolSize       int    `json:"redis_poolsize"`
+	RedisMQSize         int    `json:"redis_mqsize"`
+	MaxProcs            int    `json:"max_procs"`
+	MaxSubscriberPerKey int    `json:"max_subscriber_per_key"`
 }
 
 func InitConfig(file string) (*Config, error) {
@@ -42,22 +43,23 @@ func InitConfig(file string) (*Config, error) {
 	}
 
 	cf := &Config{
-		Addr:               "localhost",
-		Port:               8080,
-		PprofAddr:          "localhost",
-		PprofPort:          8080,
-		PubAddr:            "localhost",
-		PubPort:            8080,
-		Pprof:              1,
-		LongpollingTimeout: 300,
-		MessageTimeout:     7200,
-		Log:                "./gopush.log",
-		RedisNetwork:       "tcp",
-		RedisAddr:          "localhost:6379",
-		RedisTimeout:       28800,
-		RedisPoolSize:      50,
-		RedisMQSize:        20,
-		MaxProcs:           runtime.NumCPU(),
+		Addr:                "localhost",
+		Port:                8080,
+		PprofAddr:           "localhost",
+		PprofPort:           8080,
+		PubAddr:             "localhost",
+		PubPort:             8080,
+		Pprof:               1,
+		LongpollingTimeout:  300,
+		MessageTimeout:      7200,
+		Log:                 "./gopush.log",
+		RedisNetwork:        "tcp",
+		RedisAddr:           "localhost:6379",
+		RedisTimeout:        28800,
+		RedisPoolSize:       50,
+		RedisMQSize:         20,
+		MaxSubscriberPerKey: 0, // no limit
+		MaxProcs:            runtime.NumCPU(),
 	}
 	if err = json.Unmarshal(c, cf); err != nil {
 		Log.Printf("json.Unmarshal(\"%s\", cf) failed (%s)", string(c), err.Error())
diff --git a/gopush.conf b/gopush.conf
@@ -14,5 +14,6 @@
   "redis_timeout": 28800,
   "redis_poolsize": 50,
   "redis_mqsize": 20,
-  "max_procs": 4
+  "max_procs": 4,
+  "max_subscriber_per_key": 3
 }
diff --git a/job.go b/job.go
@@ -12,16 +12,10 @@ var (
 func CleanConnKeyJob() {
 	for {
 		key := <-ConnectedKeyCh
-
-		for {
-			if err := RedisHDel(ConnectedKey, key); err != nil {
-				Log.Printf("RedisHDel(\"%s\", \"%s\") failed (%s)", ConnectedKey, key, err.Error())
-				// if failed, sleep 1 second and retry
-				time.Sleep(1 * time.Second)
-				continue
-			}
-
-			break
+		if err := RedisDecr(ConnectedKey, key); err != nil {
+			Log.Printf("RedisDecrBy(\"%s\", \"%s\") failed (%s)", ConnectedKey, key, err.Error())
+			// if failed, sleep 1 second and retry
+			time.Sleep(1 * time.Second)
 		}
 	}
 }
diff --git a/main.go b/main.go
@@ -45,8 +45,6 @@ func main() {
 	runtime.GOMAXPROCS(Conf.MaxProcs)
 	// init redis
 	InitRedis()
-	// init clean connected key job
-	go CleanConnKeyJob()
 	// for test client
 	http.HandleFunc("/client", Client)
 	// sub
diff --git a/push.go b/push.go
@@ -112,53 +112,32 @@ func Subscribe(ws *websocket.Conn) {
 		return
 	}
 	// Generate Key
-	key = pusher.Key(key)
+	Log.Printf("Client (%s) subscribe to key %s", ws.Request().RemoteAddr, key)
 	// check multi cli connected
-	exists, err := RedisHExists(ConnectedKey, key)
-	if err != nil {
-		Log.Printf("RedisHExists(\"%s\", \"%s\") failed (%s)", ConnectedKey, key, err.Error())
-		if err = responseWriter(ws, InternalErr, result); err != nil {
-			Log.Printf("responseWriter failed (%s)", err.Error())
-		}
-
-		return
-	}
+	if Conf.MaxSubscriberPerKey > 0 {
+		ok, err := multiCliCheck(key)
+		if err != nil {
+			if err = responseWriter(ws, InternalErr, result); err != nil {
+				Log.Printf("responseWriter failed (%s)", err.Error())
+			}
 
-	if exists == 1 {
-		Log.Printf("key : %s already has a client sub to", key)
-		if err = responseWriter(ws, MultiCliErr, result); err != nil {
-			Log.Printf("responseWriter failed (%s)", err.Error())
+			return
 		}
 
-		return
-	}
-	// set connected flag
-	cliNum, err := RedisHSetnx(ConnectedKey, key, "")
-	if err != nil {
-		Log.Printf("RedisHSetnx(\"%s\", \"%s\", \"\") failed (%s)", ConnectedKey, key, err.Error())
-		if err = responseWriter(ws, InternalErr, result); err != nil {
-			Log.Printf("responseWriter failed (%s)", err.Error())
-		}
+		defer func() {
+			if err = RedisDecr(ConnectedKey, key); err != nil {
+				Log.Printf("RedisDecr(\"%s\", \"%s\") failed (%s)", ConnectedKey, key, err.Error())
+			}
+		}()
 
-		return
-	}
+		if !ok {
+			if err = responseWriter(ws, MultiCliErr, result); err != nil {
+				Log.Printf("responseWriter failed (%s)", err.Error())
+			}
 
-	if cliNum == 0 {
-		Log.Printf("key : %s already has a client sub to", key)
-		if err = responseWriter(ws, MultiCliErr, result); err != nil {
-			Log.Printf("responseWriter failed (%s)", err.Error())
+			return
 		}
-
-		return
 	}
-
-	defer func() {
-		if err := RedisHDel(ConnectedKey, key); err != nil {
-			Log.Printf("RedisHDel(\"%s\", \"%s\") failed (%s)", ConnectedKey, key, err.Error())
-			// retry to delete the key
-			ConnectedKeyCh <- key
-		}
-	}()
 	// redis routine for receive pub message or error
 	redisC, psc, err := RedisSub(key)
 	if err != nil {
@@ -207,6 +186,22 @@ func Subscribe(ws *websocket.Conn) {
 	}
 }
 
+func multiCliCheck(key string) (bool, error) {
+	// incr client num
+	cliNum, err := RedisIncr(ConnectedKey, key)
+	if err != nil {
+		Log.Printf("RedisIncr(\"%s\", \"%s\") failed (%s)", ConnectedKey, key, err.Error())
+		return false, err
+	}
+	// check
+	if cliNum > Conf.MaxSubscriberPerKey {
+		Log.Printf("key %s has %d subscribers exceed %d", key, cliNum, Conf.MaxSubscriberPerKey)
+		return false, nil
+	}
+
+	return true, nil
+}
+
 func netRead(ws *websocket.Conn) chan error {
 	c := make(chan error, 1)
 	// client close or network error, go routine exit
diff --git a/redis.go b/redis.go
@@ -32,6 +32,11 @@ func InitRedis() {
 			return c, err
 		},
 	}
+
+	if Conf.MaxSubscriberPerKey > 0 {
+		// init clean connected key job
+		go CleanConnKeyJob()
+	}
 }
 
 func RedisUnSub(key string, psc redis.PubSubConn) error {
@@ -300,6 +305,59 @@ func RedisHExists(key, field string) (int, error) {
 	return v, nil
 }
 
+func RedisHGet(key, field string) (int, error) {
+	c := redisPool.Get()
+	defer c.Close()
+	reply, err := c.Do("HGET", key, field)
+	if err != nil {
+		Log.Printf("c.Do(\"HGET\", \"%s\") failed (%s)", key, err.Error())
+		return 0, err
+	}
+
+	if reply == nil {
+		return -1, nil
+	}
+
+	v, err := redis.Int(reply, nil)
+	if err != nil {
+		Log.Printf("redis.Int() failed (%s)", err.Error())
+		return 0, err
+	}
+
+	return v, nil
+}
+
+func RedisIncr(key, field string) (int, error) {
+	c := redisPool.Get()
+	defer c.Close()
+	reply, err := c.Do("HINCRBY", key, field, 1)
+	if err != nil {
+		Log.Printf("c.Do(\"HINCRBY\", \"%s\", 1) failed (%s)", key, err.Error())
+		return 0, err
+	}
+
+	v, err := redis.Int(reply, nil)
+	if err != nil {
+		Log.Printf("redis.Int() failed (%s)", err.Error())
+		return 0, err
+	}
+
+	return v, nil
+}
+
+func RedisDecr(key, field string) error {
+	c := redisPool.Get()
+	defer c.Close()
+	_, err := c.Do("HINCRBY", key, field, -1)
+	if err != nil {
+		Log.Printf("c.Do(\"HINCRBY\", \"%s\", -1) failed (%s)", key, err.Error())
+		ConnectedKeyCh <- key
+		return err
+	}
+
+	return nil
+}
+
 func RedisHDel(key, field string) error {
 	c := redisPool.Get()
 	defer c.Close()

Original file line number	Diff line number	Diff line change
`@@ -14,5 +14,6 @@`
`14`	`14`	`"redis_timeout": 28800,`
`15`	`15`	`"redis_poolsize": 50,`
`16`	`16`	`"redis_mqsize": 20,`
`17`		`- "max_procs": 4`
	`17`	`+ "max_procs": 4,`
	`18`	`+ "max_subscriber_per_key": 3`
`18`	`19`	`}`