探索Golang Redis实现发布订阅功能实例

type Pubhub struct {

	// 自定义实现的map
	dataDict dict.ConcurrentDict

	// 该锁的颗粒度太大
	//locker sync.RWMutex

	locker *locker.Locker // 自定义一个分布锁
}

// SUBSCRIBE channel [channel ...]
func (p *Pubhub) Subscribe(c abstract.Connection, args [][]byte) protocol.Reply {
	iflen(args) < 1 {
		return protocol.NewArgNumErrReply("subscribe")
	}
	// 通道名
	keys := make([]string, 0, len(args))
	for _, arg := range args {
		keys = append(keys, string(arg))
	}
	// 加锁
	p.locker.Locks(keys...)
	defer p.locker.Unlocks(keys...)
	for _, arg := range args {
		chanName := string(arg)
		// 记录当前客户端连接订阅的通道
		c.Subscribe(chanName)
		// 双向链表，记录通道下的客户端连接
		var l *list.LinkedList
		raw, exist := p.dataDict.Get(chanName)
		if !exist { // 说明该channel第一次使用
			l = list.NewLinkedList()
			p.dataDict.Put(chanName, l)
		} else {
			l, _ = raw.(*list.LinkedList)
		}
		// 未订阅
		if !l.Contain(func(actual interface{}) bool {
			return c == actual
		}) {
			// 如果不重复，那就记录订阅
			logger.Debug("subscribe channel [" + chanName + "] success")
			l.Add(c)
		}
		// 回复客户端消息
		_, err := c.Write(channelMsg(_subscribe, chanName, c.SubCount()))
		if err != nil {
			logger.Warn(err)
		}
	}
	return protocol.NewNoReply()
}

// 取消订阅
// unsubscribes itself from all the channels using the UNSUBSCRIBE command without additional arguments
func (p *Pubhub) Unsubscribe(c abstract.Connection, args [][]byte) protocol.Reply {
	var channels []string
	iflen(args) < 1 { // 取消全部
		channels = c.GetChannels()
	} else { // 取消指定channel
		channels = make([]string, len(args))
		for i, v := range args {
			channels[i] = string(v)
		}
	}
	p.locker.Locks(channels...)
	defer p.locker.Unlocks(channels...)
	// 说明已经没有订阅的通道
	iflen(channels) == 0 {
		c.Write(noChannelMsg())
	}
	for _, channel := range channels {
		// 从客户端中删除当前通道
		c.Unsubscribe(channel)
		// 获取链表
		raw, ok := p.dataDict.Get(channel)
		if ok {
			// 从链表中删除当前客户端
			l, _ := raw.(*list.LinkedList)
			l.DelAllByVal(func(actual interface{}) bool {
				return c == actual
			})
			// 如果链表为空，清理map
			if l.Len() == 0 {
				p.dataDict.Delete(channel)
			}
		}
		c.Write(channelMsg(_unsubscribe, channel, c.SubCount()))
	}
	return protocol.NewNoReply()
}

func (p *Pubhub) Publish(self abstract.Connection, args [][]byte) protocol.Reply {
	iflen(args) != 2 {
		return protocol.NewArgNumErrReply("publish")
	}
	channelName := string(args[0])
	// 加锁
	p.locker.Locks(channelName)
	defer p.locker.Unlocks(channelName)
	raw, ok := p.dataDict.Get(channelName)
	if ok {
		var sendSuccess int64
		var failedClient = make(map[interface{}]struct{})
		// 取出链表
		l, _ := raw.(*list.LinkedList)
		// 遍历链表
		l.ForEach(func(i int, val interface{}) bool {
			conn, _ := val.(abstract.Connection)
			if conn.IsClosed() {
				failedClient[val] = struct{}{}
				returntrue
			}
			if val == self { //不给自己发送
				returntrue
			}
			// 发送数据
			conn.Write(publisMsg(channelName, string(args[1])))
			sendSuccess++
			returntrue
		})
		// 剔除客户端
		iflen(failedClient) > 0 {
			removed := l.DelAllByVal(func(actual interface{}) bool {
				_, ok := failedClient[actual]
				return ok
			})
			logger.Debugf("del %d closed client", removed)
		}
		// 返回发送的客户端数量
		return protocol.NewIntegerReply(sendSuccess)
	}
	// 如果channel不存在
	return protocol.NewIntegerReply(0)
}

type Pubhub struct {
	// 自定义实现的map
	dataDict dict.ConcurrentDict
	// 该锁的颗粒度太大
	//locker sync.RWMutex
	locker *locker.Locker // 自定义一个分布锁
}

type Locker struct {
	mu []*sync.RWMutex
	mask uint32
}
// 顺序加锁(互斥)
func (l *Locker) Locks(keys ...string) {
	indexs := l.toLockIndex(keys...)
	for _, index := range indexs {
		mu := l.mu[index]
		mu.Lock()
	}
}
// 顺序解锁(互斥)
func (l *Locker) Unlocks(keys ...string) {
	indexs := l.toLockIndex(keys...)
	for _, index := range indexs {
		mu := l.mu[index]
		mu.Unlock()
	}
}
func (l *Locker) toLockIndex(keys ...string) []uint32 {
	// 将key转成 切片索引[0,mask]
	mapIndex := make(map[uint32]struct{}) // 去重
	for _, key := range keys {
		mapIndex[l.spread(utils.Fnv32(key))] = struct{}{}
	}
	indices := make([]uint32, 0, len(mapIndex))
	for k := range mapIndex {
		indices = append(indices, k)
	}
	// 对索引排序
	sort.Slice(indices, func(i, j int) bool {
		return indices[i] < indices[j]
	})
	return indices
}