token-count.py (1532B)
1 """ 2 This is a WebTransport handler that reads tokens sent from client streams and 3 tracks how many times the token was sent. When a client sends a token, the 4 server will send the total token count back to the client. 5 """ 6 streams_dict = {} 7 8 def session_established(session): 9 # When a WebTransport session is established, a bidirectional stream is 10 # created by the server, which is used to echo back stream data from the 11 # client. 12 session.create_bidirectional_stream() 13 14 15 def stream_data_received(session, 16 stream_id: int, 17 data: bytes, 18 stream_ended: bool): 19 count = session.stash.take(data) or 0 20 count += 1 21 session.stash.put(key=data, value=count) 22 # If a stream is unidirectional, create a new unidirectional stream and echo 23 # the token count on that stream. 24 if session.stream_is_unidirectional(stream_id): 25 if (session.session_id, stream_id) not in streams_dict.keys(): 26 new_stream_id = session.create_unidirectional_stream() 27 streams_dict[(session.session_id, stream_id)] = new_stream_id 28 session.send_stream_data(streams_dict[(session.session_id, stream_id)], 29 str(count).encode()) 30 if (stream_ended): 31 del streams_dict[(session.session_id, stream_id)] 32 return 33 # Otherwise (e.g. if the stream is bidirectional), echo back the token count 34 # on the same stream. 35 session.send_stream_data(stream_id, str(count).encode())