/src/uWebSockets/src/TopicTree.h

Source (jump to first uncovered line)
/*
 * Authored by Alex Hultman, 2018-2021.
 * Intellectual property of third-party.

 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at

 *     http://www.apache.org/licenses/LICENSE-2.0

 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef UWS_TOPICTREE_H
#define UWS_TOPICTREE_H

#include <map>
#include <list>
#include <iostream>
#include <unordered_set>
#include <utility>
#include <memory>
#include <unordered_map>
#include <vector>
#include <string_view>
#include <functional>
#include <set>
#include <string>
#include <exception>

namespace uWS {

struct Subscriber;

struct Topic : std::unordered_set<Subscriber *> {

    Topic(std::string_view topic) : name(topic) {

    }

    std::string name;
};

struct Subscriber {

    template <typename, typename> friend struct TopicTree;

private:
    /* We use a factory */
    Subscriber() = default;

    /* State of prev, next does not matter unless we are needsDrainage() since we are not in the list */
    Subscriber *prev, *next;

    /* Any one subscriber can be part of at most 32 publishes before it needs a drain,
     * or whatever encoding of runs or whatever we might do in the future */
    uint16_t messageIndices[32];

    /* This one matters the most, if it is 0 we are not in the list of drainableSubscribers */
    unsigned char numMessageIndices = 0;

public:

    /* We have a list of topics we subscribe to (read by WebSocket::iterateTopics) */
    std::set<Topic *> topics;

    /* User data */
    void *user;

    bool needsDrainage() {
        return numMessageIndices;
    }
};

template <typename T, typename B>
struct TopicTree {

    enum IteratorFlags {
        LAST = 1,
        FIRST = 2
    };

    /* Whomever is iterating this topic is locked to not modify its own list */
    Subscriber *iteratingSubscriber = nullptr;

private:

    /* The drain callback must not publish, unsubscribe or subscribe.
     * It must only cork, uncork, send, write */
    std::function<bool(Subscriber *, T &, IteratorFlags)> cb;

    /* The topics */
    std::unordered_map<std::string_view, std::unique_ptr<Topic>> topics;

    /* List of subscribers that needs drainage */
    Subscriber *drainableSubscribers = nullptr;

    /* Palette of outgoing messages, up to 64k */
    std::vector<T> outgoingMessages;

    void checkIteratingSubscriber(Subscriber *s) {
        /* Notify user that they are doing something wrong here */
        if (iteratingSubscriber == s) {
            std::cerr << "Error: WebSocket must not subscribe or unsubscribe to topics while iterating its topics!" << std::endl;
            std::terminate();
        }
    }

    /* Warning: does NOT unlink from drainableSubscribers or modify next, prev. */
    void drainImpl(Subscriber *s) {
        /* Before we call cb we need to make sure this subscriber will not report needsDrainage()
         * since WebSocket::send will call drain from within the cb in that case.*/
        int numMessageIndices = s->numMessageIndices;
        s->numMessageIndices = 0;

        /* Then we emit cb */
        for (int i = 0; i < numMessageIndices; i++) {
            T &outgoingMessage = outgoingMessages[s->messageIndices[i]];

            int flags = (i == numMessageIndices - 1) ? LAST : 0;

            /* Returning true will stop drainage short (such as when backpressure is too high) */
            if (cb(s, outgoingMessage, (IteratorFlags)(flags | (i == 0 ? FIRST : 0)))) {
                break;
            }
        }
    }

    void unlinkDrainableSubscriber(Subscriber *s) {
        if (s->prev) {
            s->prev->next = s->next;
        }
        if (s->next) {
            s->next->prev = s->prev;
        }
        /* If we are the head, then we also need to reset the head */
        if (drainableSubscribers == s) {
            drainableSubscribers = s->next;
        }
    }

public:

    TopicTree(std::function<bool(Subscriber *, T &, IteratorFlags)> cb) : cb(cb) {

    }

    /* Returns nullptr if not found */
    Topic *lookupTopic(std::string_view topic) {
        auto it = topics.find(topic);
        if (it == topics.end()) {
            return nullptr;
        }
        return it->second.get();
    }

    /* Subscribe fails if we already are subscribed */
    Topic *subscribe(Subscriber *s, std::string_view topic) {
        /* Notify user that they are doing something wrong here */
        checkIteratingSubscriber(s);

        /* Lookup or create new topic */
        Topic *topicPtr = lookupTopic(topic);
        if (!topicPtr) {
            Topic *newTopic = new Topic(topic);
            topics.insert({std::string_view(newTopic->name.data(), newTopic->name.length()), std::unique_ptr<Topic>(newTopic)});
            topicPtr = newTopic;
        }

        /* Insert us in topic, insert topic in us */
        auto [it, inserted] = s->topics.insert(topicPtr);
        if (!inserted) {
            return nullptr;
        }
        topicPtr->insert(s);

        /* Success */
        return topicPtr;
    }

    /* Returns ok, last, newCount */
    std::tuple<bool, bool, int> unsubscribe(Subscriber *s, std::string_view topic) {
        /* Notify user that they are doing something wrong here */
        checkIteratingSubscriber(s);

        /* Lookup topic */
        Topic *topicPtr = lookupTopic(topic);
        if (!topicPtr) {
            /* If the topic doesn't exist we are assumed to still be subscribers of something */
            return {false, false, -1};
        }

        /* Erase from our list first */
        if (s->topics.erase(topicPtr) == 0) {
            return {false, false, -1};
        }

        /* Remove us from topic */
        topicPtr->erase(s);

        int newCount = (int) topicPtr->size();

        /* If there is no subscriber to this topic, remove it */
        if (!topicPtr->size()) {
            /* Unique_ptr deletes the topic */
            topics.erase(topic);
        }

        /* If we don't hold any topics we are to be freed altogether */
        return {true, s->topics.size() == 0, newCount};
    }

    /* Factory function for creating a Subscriber */
    Subscriber *createSubscriber() {
        return new Subscriber();
    }

    /* This is used to end a Subscriber, before freeing it */
    void freeSubscriber(Subscriber *s) {

        /* I guess we call this one even if we are not subscribers */
        if (!s) {
            return;
        }

        /* For all topics, unsubscribe */
        for (Topic *topicPtr : s->topics) {
            /* If we are the last subscriber, simply remove the whole topic */
            if (topicPtr->size() == 1) {
                topics.erase(topicPtr->name);
            } else {
                /* Otherwise just remove us */
                topicPtr->erase(s);
            }
        }

        /* We also need to unlink us */
        if (s->needsDrainage()) {
            unlinkDrainableSubscriber(s);
        }

        delete s;
    }

    /* Mainly used by WebSocket::send to drain one socket before sending */
    void drain(Subscriber *s) {
        /* The list is undefined and cannot be touched unless needsDrainage(). */
        if (s->needsDrainage()) {
            /* This function differs from drainImpl by properly unlinking
            * the subscriber from drainableSubscribers. drainImpl does not. */
            unlinkDrainableSubscriber(s);

            /* This one always resets needsDrainage before it calls any cb's.
             * Otherwise we would stackoverflow when sending after publish but before drain. */
            drainImpl(s);
            
            /* If we drained last subscriber, also clear outgoingMessages */
            if (!drainableSubscribers) {
                outgoingMessages.clear();
            }
        }
    }

    /* Called everytime we call send, to drain published messages so to sync outgoing messages */
    void drain() {
        if (drainableSubscribers) {
            /* Drain one socket a time */
            for (Subscriber *s = drainableSubscribers; s; s = s->next) {
                /* Instead of unlinking every single subscriber, we just leave the list undefined
                 * and reset drainableSubscribers ptr below. */
                drainImpl(s);
            }
            /* Drain always clears drainableSubscribers and outgoingMessages */
            drainableSubscribers = nullptr;
            outgoingMessages.clear();
        }
    }

    /* Big messages bypass all buffering and land directly in backpressure */
    template <typename F>
    bool publishBig(Subscriber *sender, std::string_view topic, B &&bigMessage, F cb) {
        /* Do we even have this topic? */
        auto it = topics.find(topic);
        if (it == topics.end()) {
            return false;
        }

        /* For all subscribers in topic */
        for (Subscriber *s : *it->second) {

            /* If we are sender then ignore us */
            if (sender != s) {
                cb(s, bigMessage);
            }
        }

        return true;
    }

    /* Linear in number of affected subscribers */
    bool publish(Subscriber *sender, std::string_view topic, T &&message) {
        /* Do we even have this topic? */
        auto it = topics.find(topic);
        if (it == topics.end()) {
            return false;
        }

        /* If we have more than 65k messages we need to drain every socket. */
        if (outgoingMessages.size() == UINT16_MAX) {
            /* If there is a socket that is currently corked, this will be ugly as all sockets will drain
             * to their own backpressure */
            drain();
        }

        /* If nobody references this message, don't buffer it */
        bool referencedMessage = false;

        /* For all subscribers in topic */
        for (Subscriber *s : *it->second) {

            /* If we are sender then ignore us */
            if (sender != s) {

                /* At least one subscriber wants this message */
                referencedMessage = true;

                /* If we already have too many outgoing messages on this subscriber, drain it now */
                if (s->numMessageIndices == 32) {
                    /* This one does not need to check needsDrainage here but still does. */
                    drain(s);
                }

                /* Finally we can continue */
                s->messageIndices[s->numMessageIndices++] = (uint16_t)outgoingMessages.size();
                /* First message adds subscriber to list of drainable subscribers */
                if (s->numMessageIndices == 1) {
                    /* Insert us in the head of drainable subscribers */
                    s->next = drainableSubscribers;
                    s->prev = nullptr;
                    if (s->next) {
                        s->next->prev = s;
                    }
                    drainableSubscribers = s;
                }
            }
        }

        /* Push this message and return with success */
        if (referencedMessage) {
            outgoingMessages.emplace_back(message);
        }

        /* Success if someone wants it */
        return referencedMessage;
    }
};

}

#endif

Coverage Report

Created: 2025-06-13 06:09

Line	Count	Source (jump to first uncovered line)
1		/*
2		* Authored by Alex Hultman, 2018-2021.
3		* Intellectual property of third-party.
4
5		* Licensed under the Apache License, Version 2.0 (the "License");
6		* you may not use this file except in compliance with the License.
7		* You may obtain a copy of the License at
8
9		* http://www.apache.org/licenses/LICENSE-2.0
10
11		* Unless required by applicable law or agreed to in writing, software
12		* distributed under the License is distributed on an "AS IS" BASIS,
13		* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14		* See the License for the specific language governing permissions and
15		* limitations under the License.
16		*/
17
18		#ifndef UWS_TOPICTREE_H
19		#define UWS_TOPICTREE_H
20
21		#include <map>
22		#include <list>
23		#include <iostream>
24		#include <unordered_set>
25		#include <utility>
26		#include <memory>
27		#include <unordered_map>
28		#include <vector>
29		#include <string_view>
30		#include <functional>
31		#include <set>
32		#include <string>
33		#include <exception>
34
35		namespace uWS {
36
37		struct Subscriber;
38
39		struct Topic : std::unordered_set<Subscriber *> {
40
41	10.6M	Topic(std::string_view topic) : name(topic) {
42
43	10.6M	}
44
45		std::string name;
46		};
47
48		struct Subscriber {
49
50		template <typename, typename> friend struct TopicTree;
51
52		private:
53		/* We use a factory */
54	106k	Subscriber() = default;
55
56		/* State of prev, next does not matter unless we are needsDrainage() since we are not in the list */
57		Subscriber prev, next;
58
59		/* Any one subscriber can be part of at most 32 publishes before it needs a drain,
60		* or whatever encoding of runs or whatever we might do in the future */
61		uint16_t messageIndices[32];
62
63		/* This one matters the most, if it is 0 we are not in the list of drainableSubscribers */
64		unsigned char numMessageIndices = 0;
65
66		public:
67
68		/* We have a list of topics we subscribe to (read by WebSocket::iterateTopics) */
69		std::set<Topic *> topics;
70
71		/* User data */
72		void *user;
73
74	155k	bool needsDrainage() {
75	155k	return numMessageIndices;
76	155k	}
77		};
78
79		template <typename T, typename B>
80		struct TopicTree {
81
82		enum IteratorFlags {
83		LAST = 1,
84		FIRST = 2
85		};
86
87		/* Whomever is iterating this topic is locked to not modify its own list */
88		Subscriber *iteratingSubscriber = nullptr;
89
90		private:
91
92		/* The drain callback must not publish, unsubscribe or subscribe.
93		* It must only cork, uncork, send, write */
94		std::function<bool(Subscriber *, T &, IteratorFlags)> cb;
95
96		/* The topics */
97		std::unordered_map<std::string_view, std::unique_ptr<Topic>> topics;
98
99		/* List of subscribers that needs drainage */
100		Subscriber *drainableSubscribers = nullptr;
101
102		/* Palette of outgoing messages, up to 64k */
103		std::vector<T> outgoingMessages;
104
105	10.6M	void checkIteratingSubscriber(Subscriber *s) {
106		/* Notify user that they are doing something wrong here */
107	10.6M	if (iteratingSubscriber == s) {
108	0	std::cerr << "Error: WebSocket must not subscribe or unsubscribe to topics while iterating its topics!" << std::endl;
109	0	std::terminate();
110	0	}
111	10.6M	}
112
113		/* Warning: does NOT unlink from drainableSubscribers or modify next, prev. */
114	12.8k	void drainImpl(Subscriber *s) {
115		/* Before we call cb we need to make sure this subscriber will not report needsDrainage()
116		* since WebSocket::send will call drain from within the cb in that case.*/
117	12.8k	int numMessageIndices = s->numMessageIndices;
118	12.8k	s->numMessageIndices = 0;
119
120		/* Then we emit cb */
121	47.5k	for (int i = 0; i < numMessageIndices; i++) {
122	35.8k	T &outgoingMessage = outgoingMessages[s->messageIndices[i]];
123
124	35.8k	int flags = (i == numMessageIndices - 1) ? LAST : 0;
125
126		/* Returning true will stop drainage short (such as when backpressure is too high) */
127	35.8k	if (cb(s, outgoingMessage, (IteratorFlags)(flags \| (i == 0 ? FIRST : 0)))) {
128	1.14k	break;
129	1.14k	}
130	35.8k	}
131	12.8k	}
132
133	43.2k	void unlinkDrainableSubscriber(Subscriber *s) {
134	43.2k	if (s->prev) {
135	0	s->prev->next = s->next;
136	0	}
137	43.2k	if (s->next) {
138	534	s->next->prev = s->prev;
139	534	}
140		/* If we are the head, then we also need to reset the head */
141	43.2k	if (drainableSubscribers == s) {
142	43.2k	drainableSubscribers = s->next;
143	43.2k	}
144	43.2k	}
145
146		public:
147
148	5.94k	TopicTree(std::function<bool(Subscriber *, T &, IteratorFlags)> cb) : cb(cb) {
149
150	5.94k	}
151
152		/* Returns nullptr if not found */
153	10.6M	Topic *lookupTopic(std::string_view topic) {
154	10.6M	auto it = topics.find(topic);
155	10.6M	if (it == topics.end()) {
156	10.6M	return nullptr;
157	10.6M	}
158	0	return it->second.get();
159	10.6M	}
160
161		/* Subscribe fails if we already are subscribed */
162	10.6M	Topic subscribe(Subscriber s, std::string_view topic) {
163		/* Notify user that they are doing something wrong here */
164	10.6M	checkIteratingSubscriber(s);
165
166		/* Lookup or create new topic */
167	10.6M	Topic *topicPtr = lookupTopic(topic);
168	10.6M	if (!topicPtr) {
169	10.6M	Topic *newTopic = new Topic(topic);
170	10.6M	topics.insert({std::string_view(newTopic->name.data(), newTopic->name.length()), std::unique_ptr<Topic>(newTopic)});
171	10.6M	topicPtr = newTopic;
172	10.6M	}
173
174		/* Insert us in topic, insert topic in us */
175	10.6M	auto [it, inserted] = s->topics.insert(topicPtr);
176	10.6M	if (!inserted) {
177	0	return nullptr;
178	0	}
179	10.6M	topicPtr->insert(s);
180
181		/* Success */
182	10.6M	return topicPtr;
183	10.6M	}
184
185		/* Returns ok, last, newCount */
186		std::tuple<bool, bool, int> unsubscribe(Subscriber *s, std::string_view topic) {
187		/* Notify user that they are doing something wrong here */
188		checkIteratingSubscriber(s);
189
190		/* Lookup topic */
191		Topic *topicPtr = lookupTopic(topic);
192		if (!topicPtr) {
193		/* If the topic doesn't exist we are assumed to still be subscribers of something */
194		return {false, false, -1};
195		}
196
197		/* Erase from our list first */
198		if (s->topics.erase(topicPtr) == 0) {
199		return {false, false, -1};
200		}
201
202		/* Remove us from topic */
203		topicPtr->erase(s);
204
205		int newCount = (int) topicPtr->size();
206
207		/* If there is no subscriber to this topic, remove it */
208		if (!topicPtr->size()) {
209		/* Unique_ptr deletes the topic */
210		topics.erase(topic);
211		}
212
213		/* If we don't hold any topics we are to be freed altogether */
214		return {true, s->topics.size() == 0, newCount};
215		}
216
217		/* Factory function for creating a Subscriber */
218	106k	Subscriber *createSubscriber() {
219	106k	return new Subscriber();
220	106k	}
221
222		/* This is used to end a Subscriber, before freeing it */
223	106k	void freeSubscriber(Subscriber *s) {
224
225		/* I guess we call this one even if we are not subscribers */
226	106k	if (!s) {
227	0	return;
228	0	}
229
230		/* For all topics, unsubscribe */
231	10.6M	for (Topic *topicPtr : s->topics) {
232		/* If we are the last subscriber, simply remove the whole topic */
233	10.6M	if (topicPtr->size() == 1) {
234	10.6M	topics.erase(topicPtr->name);
235	10.6M	} else {
236		/* Otherwise just remove us */
237	0	topicPtr->erase(s);
238	0	}
239	10.6M	}
240
241		/* We also need to unlink us */
242	106k	if (s->needsDrainage()) {
243	39.9k	unlinkDrainableSubscriber(s);
244	39.9k	}
245
246	106k	delete s;
247	106k	}
248
249		/* Mainly used by WebSocket::send to drain one socket before sending */
250	48.8k	void drain(Subscriber *s) {
251		/* The list is undefined and cannot be touched unless needsDrainage(). */
252	48.8k	if (s->needsDrainage()) {
253		/* This function differs from drainImpl by properly unlinking
254		* the subscriber from drainableSubscribers. drainImpl does not. */
255	3.23k	unlinkDrainableSubscriber(s);
256
257		/* This one always resets needsDrainage before it calls any cb's.
258		* Otherwise we would stackoverflow when sending after publish but before drain. */
259	3.23k	drainImpl(s);
260
261		/* If we drained last subscriber, also clear outgoingMessages */
262	3.23k	if (!drainableSubscribers) {
263	2.96k	outgoingMessages.clear();
264	2.96k	}
265	3.23k	}
266	48.8k	}
267
268		/* Called everytime we call send, to drain published messages so to sync outgoing messages */
269	2.63M	void drain() {
270	2.63M	if (drainableSubscribers) {
271		/* Drain one socket a time */
272	18.9k	for (Subscriber *s = drainableSubscribers; s; s = s->next) {
273		/* Instead of unlinking every single subscriber, we just leave the list undefined
274		* and reset drainableSubscribers ptr below. */
275	9.58k	drainImpl(s);
276	9.58k	}
277		/* Drain always clears drainableSubscribers and outgoingMessages */
278	9.33k	drainableSubscribers = nullptr;
279	9.33k	outgoingMessages.clear();
280	9.33k	}
281	2.63M	}
282
283		/* Big messages bypass all buffering and land directly in backpressure */
284		template <typename F>
285	0	bool publishBig(Subscriber *sender, std::string_view topic, B &&bigMessage, F cb) {
286		/* Do we even have this topic? */
287	0	auto it = topics.find(topic);
288	0	if (it == topics.end()) {
289	0	return false;
290	0	}
291
292		/* For all subscribers in topic */
293	0	for (Subscriber s : it->second) {
294
295		/* If we are sender then ignore us */
296	0	if (sender != s) {
297	0	cb(s, bigMessage);
298	0	}
299	0	}
300
301	0	return true;
302	0	}
303
304		/* Linear in number of affected subscribers */
305	1.04M	bool publish(Subscriber *sender, std::string_view topic, T &&message) {
306		/* Do we even have this topic? */
307	1.04M	auto it = topics.find(topic);
308	1.04M	if (it == topics.end()) {
309	0	return false;
310	0	}
311
312		/* If we have more than 65k messages we need to drain every socket. */
313	1.04M	if (outgoingMessages.size() == UINT16_MAX) {
314		/* If there is a socket that is currently corked, this will be ugly as all sockets will drain
315		* to their own backpressure */
316	4	drain();
317	4	}
318
319		/* If nobody references this message, don't buffer it */
320	1.04M	bool referencedMessage = false;
321
322		/* For all subscribers in topic */
323	1.04M	for (Subscriber s : it->second) {
324
325		/* If we are sender then ignore us */
326	1.04M	if (sender != s) {
327
328		/* At least one subscriber wants this message */
329	1.04M	referencedMessage = true;
330
331		/* If we already have too many outgoing messages on this subscriber, drain it now */
332	1.04M	if (s->numMessageIndices == 32) {
333		/* This one does not need to check needsDrainage here but still does. */
334	396	drain(s);
335	396	}
336
337		/* Finally we can continue */
338	1.04M	s->messageIndices[s->numMessageIndices++] = (uint16_t)outgoingMessages.size();
339		/* First message adds subscriber to list of drainable subscribers */
340	1.04M	if (s->numMessageIndices == 1) {
341		/* Insert us in the head of drainable subscribers */
342	52.7k	s->next = drainableSubscribers;
343	52.7k	s->prev = nullptr;
344	52.7k	if (s->next) {
345	784	s->next->prev = s;
346	784	}
347	52.7k	drainableSubscribers = s;
348	52.7k	}
349	1.04M	}
350	1.04M	}
351
352		/* Push this message and return with success */
353	1.04M	if (referencedMessage) {
354	1.04M	outgoingMessages.emplace_back(message);
355	1.04M	}
356
357		/* Success if someone wants it */
358	1.04M	return referencedMessage;
359	1.04M	}
360		};
361
362		}
363
364		#endif