When working on a chat app, on any app really, it is important to take into account how the app is supposed to behave when the connection is slow or lost. The moment the user hits "send" they need to be assured that the message is being processed instead of staring at a seemingly unresponsive page. Also, in case of a connection loss, the moment it is back on all the updates that took place in the interim must be processed immediately. This is where optimistic updates and graceful reconnect come in.
First, let's talk about optimistic updates. Optimistic updates are basically when you instantly display the result of an asynchronous operation assuming it will succeed, which significantly improves perceived performance. Here's the chat app we'll be using as an example:
import { useParams } from 'react-router-dom';
import { useChat } from '../api/chat.query';
import { MessageBar } from '../../../features/messageInput';
import { ErrorPopup } from '../../../shared/ui/errorPopup/ErrorPopup';
import { ChatWindow } from './ChatWindow';
import * as S from './Chat.styles';
export const Chat = () => {
const { chatId } = useParams();
const { data, isLoading, isSuccess, error } = useChat(Number(chatId));
const { messages } = isSuccess ? data : {};
return (
error ? (
<ErrorPopup errorMessage={error.message} />
) : (
<S.Chat>
<ChatWindow
messages={messages}
isLoading={isLoading}
/>
<MessageBar />
</S.Chat>
)
)
}
We have a simple chat application where messages are requested from the server using a custom useChat hook and displayed in the ChatWindow component. The implementation of the useChat hook isn't really relevant to our optimistic updates, so we will skip it. The logic for sending messages and displaying optimistic messages lives mainly in the MessageBar and ChatWindow components. To be more specific in the useCreateMessage and useMessageStream hooks used in those components. First let's take a look at the useCreateMessage hook:
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { type NewMessage, type Message } from "../../../entities/message/model/messageSchema";
import { apiClient } from "../../../shared/config/axios.api";
type ChatCache = {
messages: Message[],
[key: string]: any,
}
export const useCreateMessage = (chatId: number) => {
const queryClient = useQueryClient();
return useMutation({
mutationFn: async (newMessage: NewMessage) => {
const response = await apiClient.post(`/chats/${chatId}/messages`, newMessage);
return response.data;
},
onMutate: async (newMessage) => {
await queryClient.cancelQueries({ queryKey: ["chat", { chatId }]});
const previousChat = queryClient.getQueryData<ChatCache>(["chat", { chatId }]);
const date = new Date().toISOString();
const optimisticMessage: Message = {
...newMessage,
id: -1,
created_at: date,
modified_at: date,
}
queryClient.setQueryData(['chat', { chatId }], (chatData?: ChatCache): ChatCache | void => {
if (!chatData) return chatData;
return {
...chatData,
messages: [
optimisticMessage,
...chatData.messages,
],
};
});
return { previousChat };
},
onError: (_err, _newMessage, context) => {
queryClient.setQueryData(
['chat', { chatId }],
context?.previousChat
);
},
retry: 3,
});
};
The useCreateMessage hook is where the actual messages are sent to the server and optimistic messages are created. First, using the mutation function we send the message to the server. Then we use the onMutate method to display the optimistic update while the request is being processed. We cancel any outgoing queries to fetch chat messages using cancelQueries to prevent the optimistic message from being overridden by stale data; save the state of the chat before the message was sent to the server using getQueryData; create an optimistic message and add it to the cache using setQueryData and finally roll back to the previous state (the state of the chat we saved in an earlier step and return from onMutate) in case an error happens using onError. Now that we know how optimistic messages get to the screen, we need to figure out how to remove them once the server is done processing the message and it's saved to the database. This is where the useMessageStream hook comes in. Let's take a look:
import { useEffect } from "react";
import { useParams } from "react-router-dom";
import { useQueryClient } from "@tanstack/react-query";
import { socket } from "../../../shared/api/realtime/socket";
import { SOCKET_EVENTS } from "../../../shared/config/socket-event";
import { type Chat } from "./chatSchema";
import { messageSchema } from "../../../entities/message/model/messageSchema";
import { type Message } from "../../../entities/message/model/messageSchema";
import { produce } from "immer";
type ChatCache = {
messages: Message[],
[key: string]: any,
}
export const useMessageStream = () => {
const queryClient = useQueryClient();
const { chatId: unformattedChatId } = useParams();
const chatId = Number(unformattedChatId);
useEffect(() => {
const handleMessageReceived = (data: unknown) => {
const message = messageSchema.parse(data);
queryClient.setQueryData(['chat', { chatId }], (chatData?: ChatCache): ChatCache | void => {
if (!chatData) return chatData;
const updatedChatData = produce(chatData, draftChatData => {
const idx = draftChatData.messages.findIndex(currMsg => currMsg.client_id === message.client_id);
if (idx !== -1) {
draftChatData.messages[idx] = message;
} else {
draftChatData.messages.unshift(message);
}
})
return updatedChatData
})
}
socket.on(SOCKET_EVENTS.MESSAGE_RECEIVED, handleMessageReceived);
return () => {
socket.off(SOCKET_EVENTS.MESSAGE_RECEIVED, handleMessageReceived);
}
}, [chatId]);
}
When the message is successfully saved to the database the server emits a SOCKET_EVENTS.MESSAGE_RECEIVED event with the message itself as payload to the chat it received the message from. The listener receives the payload, finds the optimistic message using client_id (unique id created in the MessageBar component upon "send" before being passed to useCreateMessage's mutate function) and swaps it out for the real one. This is basically it for optimistic updates. You show the message before it has been processed by the server assuming the operation will succeed, and depending on whether it actually succeeds or not you either replace the optimistic message with the real one or revert the chat back to its previous state. Now let's turn our attention to graceful reconnect.
Implementing graceful reconnect can be broken down into the following steps:
- Pause requests when offline
- Sequentially flush all the queued up requests
- Display the messages we missed while offline once the connection is back
The first step is handled by TanStack by default. You can change this behavior by setting a value different from online for the networkMode property of the mutation object, but the default behavior is exactly what we want in this case, so it should be left as is. The second step is sending the paused requests to the server in the order they were created. This, too, is largely handled by TanStack Query itself. The moment the connection is back on TanStack starts processing the paused requests. However, the default behavior is to process them in parallel, which is not exactly what we want here. Suppose, we have multiple messages waiting to be processed and the moment we get our connection back they are all sent to the server at the same time. This could mess up the order in which they are supposed to be stored into the database. The last created message could be processed and saved faster than the first one if it's "heavier". Therefore we can use the scope property on the mutation options object with a key, that would require all the requests sharing the same scope.id to be processed in a sequential manner.
export const useCreateMessage = (chatId: number) => {
const queryClient = useQueryClient();
return useMutation({
mutationFn: async (newMessage: NewMessage) => {
///// QUERY FUNCTION BODY
},
scope: {
id: `chat-${chatId}`,
},
onMutate: async (newMessage) => {
///// OPTIMISTIC UPDATE LOGIC
},
onError: (_err, _newMessage, context) => {
///// ROLLBACK LOGIC
},
retry: 3,
});
};
Now, Step 3, is where things get interesting. While the user was offline, they could've missed some messages. In order to bring the user's chat window state up to date we could invalidate the query key used in the useChat hook to refetch the updated message history. However, this way along with fetching the new messages we would also be refetching data that hasn't changed, resulting in wasted bandwidth. What we could do instead is only request messages that were missed and update the cache accordingly. For this we're going to use a custom useMissedMessages hook. Here's what it looks like:
import { useEffect } from "react";
import { useParams } from "react-router-dom";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { type Message } from "../../../entities/message/model/messageSchema";
import { apiClient } from "../../../shared/config/axios.api";
import { messageListSchema } from "../../../entities/message/model/messageSchema";
import { ZodError } from "zod";
import { socket } from "../../../shared/api/realtime/socket";
import axios from "axios";
type ChatCache = {
messages: Message[],
[key: string]: any,
}
export const useMissedMessages = () => {
const queryClient = useQueryClient();
const { chatId: unformattedChatId } = useParams();
const chatId = Number(unformattedChatId);
const { mutateAsync } = useMutation({
mutationFn: async (syncId: number) => {
if (!syncId) return;
try {
const response = await apiClient.get(`/chats/${chatId}/messages/missed`, {
params: {
after: syncId,
}
});
const parsedData = messageListSchema.parse(response.data);
return parsedData;
} catch (error) {
if (axios.isAxiosError(error)) {
const responseData = error.response?.data;
let errorMessage = responseData ? responseData.error : error.message;
throw new Error(errorMessage);
}
if (error instanceof ZodError || error instanceof Error) {
throw new Error(error.message);
}
throw error;
}
},
});
const handleConnect = async () => {
const currChatData = queryClient.getQueryData<ChatCache>(['chat', { chatId }]);
if (!currChatData || currChatData.messages.length === 0) {
return queryClient.invalidateQueries({ queryKey: ['chat', { chatId }] });
}
let syncId = null;
for (const { id } of currChatData.messages) {
if (id !== -1) {
syncId = id;
break;
}
}
if (!syncId) return;
const data = await mutateAsync(syncId);
queryClient.setQueryData<ChatCache>(['chat', { chatId }], chatData => {
if (!chatData || !data) return chatData;
const { messages } = chatData;
const messagesMap = new Map<string, Message>();
for (const msg of messages) {
messagesMap.set(msg.client_id, msg);
}
for (const msg of data) {
messagesMap.set(msg.client_id, msg);
}
const updatedMessages = [...messagesMap.values()].sort(
(a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime()
);
return {
...chatData,
messages: updatedMessages,
}
})
}
useEffect(() => {
socket.on('connect', handleConnect);
return () => {
socket.off('connect', handleConnect);
}
}, [chatId]);
}
Although normally we would use useQuery or useSuspenseQuery to fetch data, we're using useMutation here to request data imperatively. We want to request missed messages only when the connection is established after a connection loss, not every time the hook runs or the component rerenders. Inside handleConnect we search for the id of the last successfully created message in the chat before the connection dropped (syncId), request all subsequent messages and update the cache with the new values. If the chat is empty we fallback to simply refetching the entire chat history, which in this scenario is justified. There is one thing worth noting. The request for missed messages and the queued up messages are sent to the server at the same time. If one of the queued messages was saved to the database before the request for the missed messages hit the server, the newly saved message would be sent back to the client along with the missed messages resulting in a race condition. This is why we use a map to avoid accidental duplicates on the client side.
And that is about it! As you can see, chat apps aren't merely about sending and receiving messages over a websocket connection. There are many edge cases and UX concerns that need to be kept in mind when designing chat systems. TanStack Query is an excellent tool that significantly simplifies the development process, but as developers a large portion of the heavy-lifting still falls on our shoulders. It's up to us to make sure a sudden connection drop or slow internet doesn't ruin user experience or at least give users a clear indication of what is going on, instead of having them wonder whether their phone is freezing, the app is glitching, they accidentally cleared the input field or one of a million other things.
Top comments (0)