Repository: DiceDB/dice
File: internal/querywatcher/query_watcher.go (Go)

Code before:
	_, err := syscall.Write(clientFD, encodedResult)
	if err != nil {
		w.removeWatcher(query, clientFD)
	}
}

// serveAdhocQueries listens for adhoc queries, executes them, and sends the result back to the client.
func (w *QWatchManager) serveAdhocQueries(ctx context.Context) {
	for {
		select {
		case query := <-AdhocQueryChan:
			result, err := w.runQuery(&query.Query)
			query.ResponseChan <- AdhocQueryResult{
				Result:      result,
				Fingerprint: generateQueryFingerprint(&query.Query),
				Err:         err,
			}
		case <-ctx.Done():
			return
		}
	}
}

// generateQueryFingerprint used to creating a common fingerprint for similar queries
func generateQueryFingerprint(query *DSQLQuery) string {
	hash := sha256.Sum256([]byte(query.KeyRegex))
	fingerprint := hex.EncodeToString(hash[:])
	return "f_" + fingerprint
}

// addWatcher adds a client as a watcher to a query.
func (w *QWatchManager) addWatcher(query *DSQLQuery, clientFD int, cacheChan chan *[]dstore.KeyValue) {
	queryFingerprint := generateQueryFingerprint(query)
	queriesMap, ok := w.FingerPrintCache.Get(queryFingerprint)
	if ok {
		// Fingerprint exists, update the existing map with the new query
		queriesMap[clientFD] = query
		w.FingerPrintCache.Put(queryFingerprint, queriesMap)
	} else {
		// Fingerprint doesn't exist, create a new map with the query
		newMap := make(map[int]*DSQLQuery)
		newMap[clientFD] = query
		w.FingerPrintCache.Put(queryFingerprint, newMap)
	}

	clients, _ := w.WatchList.LoadOrStore(queryFingerprint, &sync.Map{})
	clients.(*sync.Map).Store(clientFD, struct{}{})

	w.QueryCacheMu.Lock()
	defer w.QueryCacheMu.Unlock()


Review comment (security):
I don't think sha256 is a good choice here since we don't care about cryptographic security. We want something fast, maybe something like FarmHash?

Code after:
			result, err := w.runQuery(&query.Query)
			query.ResponseChan <- AdhocQueryResult{
				Result:      result,
				Fingerprint: query.Query.Fingerprint,
				Err:         err,
			}
		case <-ctx.Done():
			return
		}
	}
}

// addWatcher adds a client as a watcher to a query.
func (w *QueryManager) addWatcher(query *DSQLQuery, clientFD int, cacheChan chan *[]dstore.KeyValue) {
	queryString := query.String()

	clients, _ := w.WatchList.LoadOrStore(queryString, &sync.Map{})
	clients.(*sync.Map).Store(clientFD, struct{}{})

	w.QueryCacheMu.Lock()
	defer w.QueryCacheMu.Unlock()

	cache := newCacheStore()
	// Hydrate the cache with data from all shards.
	// TODO: We need to ensure we receive cache data from all shards once we have multithreading in place.
	//  For now we only expect one update.
	kvs := <-cacheChan
	for _, kv := range *kvs {
		((*swiss.Map[string, *dstore.Obj])(cache)).Put(kv.Key, kv.Value)
	}

	w.QueryCache.Put(query.Fingerprint, cache)
}

// removeWatcher removes a client from the watchlist for a query.
func (w *QueryManager) removeWatcher(query *DSQLQuery, clientFD int) {
	queryString := query.String()
	if clients, ok := w.WatchList.Load(queryString); ok {
		clients.(*sync.Map).Delete(clientFD)
		log.Info(fmt.Sprintf("client '%d' no longer watching query: %s", clientFD, queryString))

		// If no more clients for this query, remove the query from WatchList
		if w.clientCount(clients.(*sync.Map)) == 0 {
			w.WatchList.Delete(queryString)

			// Remove this Query's cached data.
			w.QueryCacheMu.Lock()
			w.QueryCache.Delete(query.Fingerprint)
			w.QueryCacheMu.Unlock()

			log.Info(fmt.Sprintf("no longer watching query: %s", queryString))