linux网络编程十六：I/O复用的应用-poll简单实现聊天室程序

这次我们以poll为例实现一个简单的聊天室程序，实现让所有用户同时在线群聊，分为客户端和服务端两部分。

客户端有两个功能：一是从标准输入终端读取用数据，并将用户数据发关至服务器；二是往标准输出终端打印服务器发送给它的数据。

服务端功能是接收客户数据，并把客户数据发送给每一个登录到该服务器上的客户端（数据发送者除外）。


1. 客户端使用poll监听用户输入和网络连接，并利用splice函数将用户输入内容直接定向到网络连接上发关，实现零拷贝，提高效率。

2. 服务端使用poll同时管理监听socket和连接socket，并使用牺牲空间来换取时间的策略来提高服务器性能。

3. 客户端代码

//客户端
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <assert.h>
#include <poll.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define BUFFER_SIZE 64

int main(int argc, char **argv)
{
	if (argc != 3) {
		fprintf(stderr, "Usage: %s ip port\n", basename(argv[0]));
		return 1;
	}
	
	const char *ip = argv[1];
	int port = atoi(argv[2]);
	
	struct sockaddr_in server_address;
	bzero(&server_address, sizeof(server_address));
	server_address.sin_family = AF_INET;
	server_address.sin_port = htons(port);
	inet_pton(AF_INET, ip, &server_address.sin_addr);
	
	int sockfd = socket(PF_INET, SOCK_STREAM, 0);
	assert(sockfd >= 0);
	
	if (connect(sockfd, (struct sockaddr*)&server_address, sizeof(server_address)) < 0) {
		printf("connection failed\n");
		close(sockfd);
		return 1;
	}
	
	pollfd fds[2];
	fds[0].fd = 0;
	fds[0].events = POLLIN;
	fds[0].revents = 0;
	fds[1].fd = sockfd;
	fds[1].events = POLLIN | POLLRDHUP;
	fds[1].revents = 0;
	
	char read_buf[BUFFER_SIZE];
	int pipefd[2];
	
	int ret = pipe(pipefd);
	assert(ret != -1);
	
	while (1) {
		ret = poll(fds, 2, -1);
		if (ret < 0) {
			printf("poll failed\n");
			break;
		}
		
		if (fds[1].revents & POLLRDHUP) {
			printf("server close the connection\n");
			break;
		}
		else if (fds[1].revents & POLLIN) {
			memset(read_buf, '\0', BUFFER_SIZE);
			recv(fds[1].fd, read_buf, BUFFER_SIZE-1, 0);
			printf("%s\n", read_buf);
		}
		
		if (fds[0].revents & POLLIN) {
			ret = splice(0, NULL, pipefd[1], NULL, 32768, SPLICE_F_MORE | SPLICE_F_MOVE);
			ret = splice(pipefd[0], NULL, sockfd, NULL, 32768, SPLICE_F_MORE | SPLICE_F_MOVE);
		}
		
	}
	
	close(sockfd);
		
	
	
	return 0;
}

4. 服务端代码

//服务端
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <fcntl.h>
#include <poll.h>
#include <assert.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <signal.h>

#define USER_LIMIT 512		//最大用户数量
#define BUFFER_SIZE 64		//读缓冲区的大小
#define FD_LIMIT 65535		//文件描述符数量限制

struct client_data
{
	sockaddr_in address;		//客户端地址
	char *write_buf;			//发送到客户端的缓存区
	char buf[BUFFER_SIZE];	//客户端接收的缓存区
};

//设置非阻塞
int setnonblocking(int fd)
{
	int old_option = fcntl(fd, F_GETFL);
	int new_option = old_option | O_NONBLOCK;
	fcntl(fd, F_SETFL, new_option);
	return old_option;
}

//忽略SIGPIPE信号
static int ignore_sigpipe()
{
	struct sigaction act;
	
	if (sigaction(SIGPIPE, (struct sigaction*)NULL, &act) == -1)
		return -1;
		
	if (act.sa_handler == SIG_DFL) {
		act.sa_handler = SIG_IGN;
		if (sigaction(SIGPIPE, &act, (struct sigaction*)NULL) == -1)
			return -1;
	}
	
	return 0;
}

int main(int argc, char **argv)
{
	if (argc != 2) {
		fprintf(stderr, "Usage: %s ip\n", basename(argv[0]));
		return 1;
	}
	
	if (ignore_sigpipe() == -1)
		return 1;
	
	int port = atoi(argv[1]);	
	int ret = 0;
	int error;
	
	struct sockaddr_in address;
	bzero(&address, sizeof(address));
	address.sin_family = AF_INET;
	address.sin_port = htons(port);
	address.sin_addr.s_addr = htonl(INADDR_ANY);
	
	int sockfd = socket(PF_INET, SOCK_STREAM, 0);	
	if (sockfd == -1)
		return 1;
		
	printf("server start...\n");
	
	//设置地址可重用
	int reuse = 1;
	ret = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse));
	if (ret == -1) {
		error = errno;
		while ((close(sockfd) == -1) && (errno == EINTR));
		errno = error;
		return 1;
	}
	
	printf("server reuseaddr success\n");
		
	if ((bind(sockfd, (struct sockaddr*)&address, sizeof(address)) == -1) ||
		(listen(sockfd, 5) == -1)) {
		error = errno;
		while ((close(sockfd) == -1) && (errno == EINTR));
		errno = error;
		return 1;
	}
	
	printf("server bind and listen success\n");
	
	//客户端数组，每个socket连接获得一个这样的对象，socket值作为下标索引。
	client_data *users = new client_data[FD_LIMIT];
	
	//虽然已分配了足够多的client_data对象，但为了提高poll性能，仍然有必要限制用户数量
	pollfd fds[USER_LIMIT + 1];
	
	//当前用户连接数
	int user_counter = 0;
	
	for (int i = 0; i <= USER_LIMIT; ++i) {
		fds[i].fd = -1;
		fds[i].events = 0;
	}
	
	fds[0].fd = sockfd;
	fds[0].events = POLLIN | POLLERR;
	fds[0].revents = 0;
	
	while (1) {
		ret = poll(fds, user_counter+1, -1);
		if (ret < 0) {
			fprintf(stderr, "poll failed\n");
			break;
		}
		
		for (int i = 0; i < user_counter+1; ++i) {
			if ((fds[i].fd == sockfd) && fds[i].revents & POLLIN ) {
					struct sockaddr_in client_address;
					socklen_t client_addrlength = sizeof(client_address);
					
					int connfd;
					while (((connfd = accept(sockfd, (struct sockaddr*)&client_address, &client_addrlength)) == -1) && (connfd == EINTR));
					
					if (connfd < 0) {
						fprintf(stderr, "accept failed: %s\n", strerror(errno));
						continue;
					}
					
					//如果请求太多，关闭新到的连接
					if (user_counter >= USER_LIMIT) {
						const char *info = "too many users, please later\n";
						printf("%s", info);
						
						send(connfd, info, strlen(info), 0);
						close(connfd);
						continue;
					}
					
					user_counter++;
					
					//添加新连接的客户端到用户数组
					users[connfd].address = client_address;
					setnonblocking(connfd);
					
					//将新连接的客户端加入轮询
					fds[user_counter].fd = connfd;
					fds[user_counter].events = POLLIN | POLLRDHUP | POLLERR;
					fds[user_counter].revents = 0;
					
					printf("comes a new user: %d, now have %d users\n", connfd, user_counter);
			}
			else if (fds[i].revents & POLLRDHUP) {
				//客户端关闭连接，则服务端也关闭对应的连接，并将用户总数减1
				//将最后一个连接数据复制到当前位置，删除原数据
				users[fds[i].fd] = users[fds[user_counter].fd];
				
				close(fds[i].fd);
				fds[i] = fds[user_counter];
				i--;
				user_counter--;
				printf("a client left\n");
			}
			else if (fds[i].revents & POLLIN) {
				int connfd = fds[i].fd;
				memset(users[connfd].buf, '\0', BUFFER_SIZE);
				
				ret = recv(connfd, users[connfd].buf, BUFFER_SIZE-1, 0);
				if (ret < 0) {
					//如果读取错误，则关闭连接
					if (errno != EAGAIN) {
						fprintf(stderr, "client %d recv error: %s\n", connfd, strerror(errno));
						close(connfd);
						users[fds[i].fd] = users[fds[user_counter].fd];
						fds[i] = fds[user_counter];
						i--;
						user_counter--;
					}
				}
				else if (ret == 0) {
					fprintf(stderr, "client %d code should not come to here\n", connfd);
				}
				else {
					//如果接收到客户数据，则通知其他连接写数据
					for (int j = 1; j <= user_counter; ++j) {
						if (fds[j].fd == connfd)
							continue;
							
						fds[j].events |= ~POLLIN;
						fds[j].events |= POLLOUT;
						users[fds[j].fd].write_buf = users[connfd].buf;
					}
				}
			}
			else if (fds[i].revents & POLLOUT) {
				int connfd = fds[i].fd;
				if (!users[connfd].write_buf)
					continue;
					
				ret = send(connfd, users[connfd].write_buf, strlen(users[connfd].write_buf), 0);
				users[connfd].write_buf = NULL;
				
				//写完数据后，需要重新注册可读事件
				fds[i].events |= ~POLLOUT;
				fds[i].events |= POLLIN;
			}
		}
		
	}
	
	delete[] users;
	close(sockfd);
	
	return 0;
}

参考：《linux高性能服务器编程》