Coverage for qreward / utils / socket_keepalive.py: 100.00%

1import socket (empty)

2import sys (empty)

3from typing import List (empty)

4 

5from requests.adapters import HTTPAdapter (empty)

6 

7from qreward.types import SOCKET_OPTION (empty)

8 

9 

10def aiohttp_keepalive_socket_factory( (empty)

11 addr_info, 

12 keepalive_secs: int = 60, 

13 keepalive_interval: int = 30, 

14 keepalive_cnt: int = 3, 

15) -> socket.socket: 

16 """aiohttp 保持 keepalive 连接的 socket 工厂函数 

17 

18 Args: 

19 addr_info (tuple): socket 地址信息 

20 keepalive_secs (int, optional): 保持 keepalive 连接的时间, 默认 60 秒 

21 keepalive_interval (int, optional): 保持 keepalive 连接的间隔, 默认 30 秒 

22 keepalive_cnt (int, optional): 保持 keepalive 探测的次数, 默认 3 次 

23 

24 Returns: 

25 socket: socket 对象 

26 """ 

27 

28 family, type_, proto, _, _ = addr_info (empty)

29 

30 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) (empty)

31 sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) (empty)

32 

33 if sys.platform.startswith("linux"): (empty)

34 # Linux 

35 sock.setsockopt( (empty)

36 socket.IPPROTO_TCP, 

37 socket.TCP_KEEPIDLE, 

38 keepalive_secs, 

39 ) 

40 elif sys.platform.startswith("darwin"): (empty)

41 # Darwin 

42 sock.setsockopt( (empty)

43 socket.IPPROTO_TCP, 

44 socket.TCP_KEEPALIVE, 

45 keepalive_secs, 

46 ) 

47 

48 sock.setsockopt( (empty)

49 socket.IPPROTO_TCP, 

50 socket.TCP_KEEPINTVL, 

51 keepalive_interval, 

52 ) 

53 sock.setsockopt( (empty)

54 socket.IPPROTO_TCP, 

55 socket.TCP_KEEPCNT, 

56 keepalive_cnt, 

57 ) 

58 return sock (empty)

59 

60 

61def httpx_keepalive_socket( (empty)

62 keepalive_secs: int = 60, 

63 keepalive_interval: int = 30, 

64 keepalive_cnt: int = 3, 

65) -> List[SOCKET_OPTION]: 

66 """ 

67 httpx 保持 keepalive 连接的 socket 选项 

68 

69 :param keepalive_secs: 保持 keepalive 连接的时间, 默认 60 秒 

70 :param keepalive_interval: 保持 keepalive 连接的间隔, 默认 30 秒 

71 :param keepalive_cnt: 保持 keepalive 探测的次数, 默认 3 次 

72 :return: socket 选项列表 

73 """ 

74 

75 socket_options = [ (empty)

76 (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), # 启用keepalive 

77 ] 

78 

79 if sys.platform.startswith("linux"): (empty)

80 socket_options.append( (empty)

81 ( 

82 socket.IPPROTO_TCP, 

83 socket.TCP_KEEPIDLE, 

84 keepalive_secs, 

85 ) 

86 ) 

87 elif sys.platform.startswith("darwin"): (empty)

88 socket_options.append( (empty)

89 ( 

90 socket.IPPROTO_TCP, 

91 socket.TCP_KEEPALIVE, 

92 keepalive_secs, 

93 ) 

94 ) 

95 

96 socket_options.append( (empty)

97 ( 

98 socket.IPPROTO_TCP, 

99 socket.TCP_KEEPINTVL, 

100 keepalive_interval, 

101 ) 

102 ) 

103 socket_options.append( (empty)

104 ( 

105 socket.IPPROTO_TCP, 

106 socket.TCP_KEEPCNT, 

107 keepalive_cnt, 

108 ) 

109 ) 

110 return socket_options (empty)

111 

112 

113class RequestsTCPKeepAliveAdapter(HTTPAdapter): (empty)

114 """自定义适配器：注入keepalive socket选项""" 

115 

116 def __init__(self, *args, **kwargs): (empty)

117 self._default = httpx_keepalive_socket() (empty)

118 self.socket_options = kwargs.pop("socket_options", self._default) (empty)

119 super().__init__(*args, **kwargs) (empty)

120 

121 def init_poolmanager( (empty)

122 self, 

123 connections, 

124 maxsize, 

125 block=False, 

126 **pool_kwargs, 

127 ): 

128 if self.socket_options is not None: (empty)

129 pool_kwargs["socket_options"] = self.socket_options (empty)

130 super().init_poolmanager(connections, maxsize, block, **pool_kwargs) (empty)