github.com/hdt3213/godis@v1.2.9/pubsub/pubsub.go (about)

     1  package pubsub
     2  
     3  import (
     4  	"github.com/hdt3213/godis/datastruct/list"
     5  	"github.com/hdt3213/godis/interface/redis"
     6  	"github.com/hdt3213/godis/lib/utils"
     7  	"github.com/hdt3213/godis/redis/protocol"
     8  	"strconv"
     9  )
    10  
    11  var (
    12  	_subscribe         = "subscribe"
    13  	_unsubscribe       = "unsubscribe"
    14  	messageBytes       = []byte("message")
    15  	unSubscribeNothing = []byte("*3\r\n$11\r\nunsubscribe\r\n$-1\n:0\r\n")
    16  )
    17  
    18  func makeMsg(t string, channel string, code int64) []byte {
    19  	return []byte("*3\r\n$" + strconv.FormatInt(int64(len(t)), 10) + protocol.CRLF + t + protocol.CRLF +
    20  		"$" + strconv.FormatInt(int64(len(channel)), 10) + protocol.CRLF + channel + protocol.CRLF +
    21  		":" + strconv.FormatInt(code, 10) + protocol.CRLF)
    22  }
    23  
    24  /*
    25   * invoker should lock channel
    26   * return: is new subscribed
    27   */
    28  func subscribe0(hub *Hub, channel string, client redis.Connection) bool {
    29  	client.Subscribe(channel)
    30  
    31  	// add into hub.subs
    32  	raw, ok := hub.subs.Get(channel)
    33  	var subscribers *list.LinkedList
    34  	if ok {
    35  		subscribers, _ = raw.(*list.LinkedList)
    36  	} else {
    37  		subscribers = list.Make()
    38  		hub.subs.Put(channel, subscribers)
    39  	}
    40  	if subscribers.Contains(func(a interface{}) bool {
    41  		return a == client
    42  	}) {
    43  		return false
    44  	}
    45  	subscribers.Add(client)
    46  	return true
    47  }
    48  
    49  /*
    50   * invoker should lock channel
    51   * return: is actually un-subscribe
    52   */
    53  func unsubscribe0(hub *Hub, channel string, client redis.Connection) bool {
    54  	client.UnSubscribe(channel)
    55  
    56  	// remove from hub.subs
    57  	raw, ok := hub.subs.Get(channel)
    58  	if ok {
    59  		subscribers, _ := raw.(*list.LinkedList)
    60  		subscribers.RemoveAllByVal(func(a interface{}) bool {
    61  			return utils.Equals(a, client)
    62  		})
    63  
    64  		if subscribers.Len() == 0 {
    65  			// clean
    66  			hub.subs.Remove(channel)
    67  		}
    68  		return true
    69  	}
    70  	return false
    71  }
    72  
    73  // Subscribe puts the given connection into the given channel
    74  func Subscribe(hub *Hub, c redis.Connection, args [][]byte) redis.Reply {
    75  	channels := make([]string, len(args))
    76  	for i, b := range args {
    77  		channels[i] = string(b)
    78  	}
    79  
    80  	hub.subsLocker.Locks(channels...)
    81  	defer hub.subsLocker.UnLocks(channels...)
    82  
    83  	for _, channel := range channels {
    84  		if subscribe0(hub, channel, c) {
    85  			_, _ = c.Write(makeMsg(_subscribe, channel, int64(c.SubsCount())))
    86  		}
    87  	}
    88  	return &protocol.NoReply{}
    89  }
    90  
    91  // UnsubscribeAll removes the given connection from all subscribing channel
    92  func UnsubscribeAll(hub *Hub, c redis.Connection) {
    93  	channels := c.GetChannels()
    94  
    95  	hub.subsLocker.Locks(channels...)
    96  	defer hub.subsLocker.UnLocks(channels...)
    97  
    98  	for _, channel := range channels {
    99  		unsubscribe0(hub, channel, c)
   100  	}
   101  
   102  }
   103  
   104  // UnSubscribe removes the given connection from the given channel
   105  func UnSubscribe(db *Hub, c redis.Connection, args [][]byte) redis.Reply {
   106  	var channels []string
   107  	if len(args) > 0 {
   108  		channels = make([]string, len(args))
   109  		for i, b := range args {
   110  			channels[i] = string(b)
   111  		}
   112  	} else {
   113  		channels = c.GetChannels()
   114  	}
   115  
   116  	db.subsLocker.Locks(channels...)
   117  	defer db.subsLocker.UnLocks(channels...)
   118  
   119  	if len(channels) == 0 {
   120  		_, _ = c.Write(unSubscribeNothing)
   121  		return &protocol.NoReply{}
   122  	}
   123  
   124  	for _, channel := range channels {
   125  		if unsubscribe0(db, channel, c) {
   126  			_, _ = c.Write(makeMsg(_unsubscribe, channel, int64(c.SubsCount())))
   127  		}
   128  	}
   129  	return &protocol.NoReply{}
   130  }
   131  
   132  // Publish send msg to all subscribing client
   133  func Publish(hub *Hub, args [][]byte) redis.Reply {
   134  	if len(args) != 2 {
   135  		return &protocol.ArgNumErrReply{Cmd: "publish"}
   136  	}
   137  	channel := string(args[0])
   138  	message := args[1]
   139  
   140  	hub.subsLocker.Lock(channel)
   141  	defer hub.subsLocker.UnLock(channel)
   142  
   143  	raw, ok := hub.subs.Get(channel)
   144  	if !ok {
   145  		return protocol.MakeIntReply(0)
   146  	}
   147  	subscribers, _ := raw.(*list.LinkedList)
   148  	subscribers.ForEach(func(i int, c interface{}) bool {
   149  		client, _ := c.(redis.Connection)
   150  		replyArgs := make([][]byte, 3)
   151  		replyArgs[0] = messageBytes
   152  		replyArgs[1] = []byte(channel)
   153  		replyArgs[2] = message
   154  		_, _ = client.Write(protocol.MakeMultiBulkReply(replyArgs).ToBytes())
   155  		return true
   156  	})
   157  	return protocol.MakeIntReply(int64(subscribers.Len()))
   158  }