| Michael Johnson | 89bdbcd | 2020-03-19 15:59:19 -0700 | [diff] [blame] | 1 | # Licensed under the Apache License, Version 2.0 (the "License"); you may |
| 2 | # not use this file except in compliance with the License. You may obtain |
| 3 | # a copy of the License at |
| 4 | # |
| 5 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 6 | # |
| 7 | # Unless required by applicable law or agreed to in writing, software |
| 8 | # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| 9 | # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| 10 | # License for the specific language governing permissions and limitations |
| 11 | # under the License. |
| 12 | import socket |
| 13 | |
| 14 | from requests.adapters import HTTPAdapter |
| 15 | from requests.packages.urllib3.connection import HTTPConnection |
| 16 | from requests.packages.urllib3.poolmanager import PoolManager |
| 17 | |
| 18 | |
| 19 | class SourcePortAdapter(HTTPAdapter): |
| 20 | """"Transport adapter" that allows us to set the source port.""" |
| 21 | def __init__(self, port, *args, **kwargs): |
| 22 | self._source_port = port |
| 23 | super(SourcePortAdapter, self).__init__(*args, **kwargs) |
| 24 | |
| 25 | def init_poolmanager(self, connections, maxsize, block=False): |
| 26 | # Make sure TIMED_WAIT doesn't stop us from reusing the socket |
| 27 | sock_options = HTTPConnection.default_socket_options + [ |
| 28 | (socket.SOL_SOCKET, socket.SO_REUSEADDR, 1), ] |
| 29 | self.poolmanager = PoolManager( |
| 30 | num_pools=connections, maxsize=maxsize, |
| 31 | block=block, source_address=('', self._source_port), |
| 32 | socket_options=sock_options) |