Repository: DiceDB/dice
File: internal/server/ironhawk/watch_manager.go (Go)

Code before:

	// TODO: If the key gets deleted from the database
	// delete the subscriptions against that key from all the places.
}

func (w *WatchManager) CleanupThreadWatchSubscriptions(t *IOThread) {
	w.mu.Lock()
	defer w.mu.Unlock()

	// Delete the mapping of Watch thread to client id
	delete(w.clientWatchThreadMap, t.ClientID)

	// Delete all the subscriptions of the client from the fingerprint maps
	// Note: this is an O(n) operation and hence if there are large number of clients, this might be expensive.
	// We can do a lazy deletion of the fingerprint map if this becomes a problem.
	for fp := range w.fpClientMap {
		delete(w.fpClientMap[fp], t.ClientID)
		if len(w.fpClientMap[fp]) == 0 {
			delete(w.fpClientMap, fp)
		}
	}
}

func (w *WatchManager) NotifyWatchers(c *cmd.Cmd, shardManager *shardmanager.ShardManager, t *IOThread) {
	// Use RLock instead as we are not really modifying any shared maps here.
	w.mu.RLock()
	defer w.mu.RUnlock()

	key := c.Key()
	for fp := range w.keyFPMap[key] {
		_c := w.fpCmdMap[fp]
		if _c == nil {
			// TODO: Not having a command for a fingerprint is a bug.
			continue
		}

		r, err := _c.Execute(shardManager)
		if err != nil {
			slog.Error("failed to execute command as part of watch notification",
				slog.Any("cmd", _c.String()),
				slog.Any("error", err))
			continue
		}

		for clientID := range w.fpClientMap[fp] {
			thread := w.clientWatchThreadMap[clientID]
			if thread == nil {
				// if there is no thread against the client, delete the client from the map
				delete(w.clientWatchThreadMap, clientID)
				continue
			}

Review comment (suggestion):
This change should significantly improve the through put of notifications on watches as RLocks can be acquired concurrently. The earlier code would only allow either one writer or reader to be processed at the same time.

With this change, concurrent users can read but only one writer to be processed.

Code after:
	// TODO: If the key gets deleted from the database
	// delete the subscriptions against that key from all the places.
}

func (w *WatchManager) CleanupThreadWatchSubscriptions(t *IOThread) {
	w.mu.Lock()
	defer w.mu.Unlock()

	// Delete the mapping of Watch thread to client id
	delete(w.clientWatchThreadMap, t.ClientID)

	// Delete all the subscriptions of the client from the fingerprint maps
	// Note: this is an O(n) operation and hence if there are large number of clients, this might be expensive.
	// We can do a lazy deletion of the fingerprint map if this becomes a problem.
	for fp := range w.fpClientMap {
		delete(w.fpClientMap[fp], t.ClientID)
		if len(w.fpClientMap[fp]) == 0 {
			delete(w.fpClientMap, fp)
		}
	}
}

func (w *WatchManager) NotifyWatchers(c *cmd.Cmd, shardManager *shardmanager.ShardManager, t *IOThread) {
	// Use RLock instead as we are not really modifying any shared maps here.
	w.mu.RLock()
	defer w.mu.RUnlock()

	key := c.Key()
	for fp := range w.keyFPMap[key] {
		_c := w.fpCmdMap[fp]
		if _c == nil {
			// TODO: Not having a command for a fingerprint is a bug.
			continue
		}

		r, err := _c.Execute(shardManager)
		if err != nil {
			slog.Error("failed to execute command as part of watch notification",
				slog.Any("cmd", _c.String()),
				slog.Any("error", err))
			continue
		}

		for clientID := range w.fpClientMap[fp] {
			thread := w.clientWatchThreadMap[clientID]
			if thread == nil {
				// if there is no thread against the client, delete the client from the map
				delete(w.clientWatchThreadMap, clientID)
				continue
			}