ConnectionPoolTests.swift 46 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375
  1. /*
  2. * Copyright 2021, gRPC Authors All rights reserved.
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License");
  5. * you may not use this file except in compliance with the License.
  6. * You may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * Unless required by applicable law or agreed to in writing, software
  11. * distributed under the License is distributed on an "AS IS" BASIS,
  12. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. * See the License for the specific language governing permissions and
  14. * limitations under the License.
  15. */
  16. @testable import GRPC
  17. import Logging
  18. import NIOCore
  19. import NIOEmbedded
  20. import NIOHTTP2
  21. import XCTest
  22. final class ConnectionPoolTests: GRPCTestCase {
  23. private enum TestError: Error {
  24. case noChannelExpected
  25. }
  26. private var eventLoop: EmbeddedEventLoop!
  27. private var tearDownBlocks: [() throws -> Void] = []
  28. override func setUp() {
  29. super.setUp()
  30. self.eventLoop = EmbeddedEventLoop()
  31. }
  32. override func tearDown() {
  33. XCTAssertNoThrow(try self.eventLoop.close())
  34. self.tearDownBlocks.forEach { try? $0() }
  35. super.tearDown()
  36. }
  37. private func noChannelExpected(
  38. _: ConnectionManager,
  39. _ eventLoop: EventLoop,
  40. line: UInt = #line
  41. ) -> EventLoopFuture<Channel> {
  42. XCTFail("Channel unexpectedly created", line: line)
  43. return eventLoop.makeFailedFuture(TestError.noChannelExpected)
  44. }
  45. private func makePool(
  46. waiters: Int = 1000,
  47. reservationLoadThreshold: Double = 0.9,
  48. now: @escaping () -> NIODeadline = { .now() },
  49. connectionBackoff: ConnectionBackoff = ConnectionBackoff(),
  50. delegate: GRPCConnectionPoolDelegate? = nil,
  51. onReservationReturned: @escaping (Int) -> Void = { _ in },
  52. onMaximumReservationsChange: @escaping (Int) -> Void = { _ in },
  53. channelProvider: ConnectionManagerChannelProvider
  54. ) -> ConnectionPool {
  55. return ConnectionPool(
  56. eventLoop: self.eventLoop,
  57. maxWaiters: waiters,
  58. reservationLoadThreshold: reservationLoadThreshold,
  59. assumedMaxConcurrentStreams: 100,
  60. connectionBackoff: connectionBackoff,
  61. channelProvider: channelProvider,
  62. streamLender: HookedStreamLender(
  63. onReturnStreams: onReservationReturned,
  64. onUpdateMaxAvailableStreams: onMaximumReservationsChange
  65. ),
  66. delegate: delegate,
  67. logger: self.logger.wrapped,
  68. now: now
  69. )
  70. }
  71. private func makePool(
  72. waiters: Int = 1000,
  73. delegate: GRPCConnectionPoolDelegate? = nil,
  74. makeChannel: @escaping (ConnectionManager, EventLoop) -> EventLoopFuture<Channel>
  75. ) -> ConnectionPool {
  76. return self.makePool(
  77. waiters: waiters,
  78. delegate: delegate,
  79. channelProvider: HookedChannelProvider(makeChannel)
  80. )
  81. }
  82. private func setUpPoolAndController(
  83. waiters: Int = 1000,
  84. reservationLoadThreshold: Double = 0.9,
  85. now: @escaping () -> NIODeadline = { .now() },
  86. connectionBackoff: ConnectionBackoff = ConnectionBackoff(),
  87. delegate: GRPCConnectionPoolDelegate? = nil,
  88. onReservationReturned: @escaping (Int) -> Void = { _ in },
  89. onMaximumReservationsChange: @escaping (Int) -> Void = { _ in }
  90. ) -> (ConnectionPool, ChannelController) {
  91. let controller = ChannelController()
  92. let pool = self.makePool(
  93. waiters: waiters,
  94. reservationLoadThreshold: reservationLoadThreshold,
  95. now: now,
  96. connectionBackoff: connectionBackoff,
  97. delegate: delegate,
  98. onReservationReturned: onReservationReturned,
  99. onMaximumReservationsChange: onMaximumReservationsChange,
  100. channelProvider: controller
  101. )
  102. self.tearDownBlocks.append {
  103. let shutdown = pool.shutdown()
  104. self.eventLoop.run()
  105. XCTAssertNoThrow(try shutdown.wait())
  106. controller.finish()
  107. }
  108. return (pool, controller)
  109. }
  110. func testEmptyConnectionPool() {
  111. let pool = self.makePool {
  112. self.noChannelExpected($0, $1)
  113. }
  114. XCTAssertEqual(pool.sync.connections, 0)
  115. XCTAssertEqual(pool.sync.waiters, 0)
  116. XCTAssertEqual(pool.sync.availableStreams, 0)
  117. XCTAssertEqual(pool.sync.reservedStreams, 0)
  118. pool.initialize(connections: 20)
  119. XCTAssertEqual(pool.sync.connections, 20)
  120. XCTAssertEqual(pool.sync.waiters, 0)
  121. XCTAssertEqual(pool.sync.availableStreams, 0)
  122. XCTAssertEqual(pool.sync.reservedStreams, 0)
  123. let shutdownFuture = pool.shutdown()
  124. self.eventLoop.run()
  125. XCTAssertNoThrow(try shutdownFuture.wait())
  126. }
  127. func testShutdownEmptyPool() {
  128. let pool = self.makePool {
  129. self.noChannelExpected($0, $1)
  130. }
  131. XCTAssertNoThrow(try pool.shutdown().wait())
  132. // Shutting down twice should also be fine.
  133. XCTAssertNoThrow(try pool.shutdown().wait())
  134. }
  135. func testMakeStreamWhenShutdown() {
  136. let pool = self.makePool {
  137. self.noChannelExpected($0, $1)
  138. }
  139. XCTAssertNoThrow(try pool.shutdown().wait())
  140. let stream = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  141. $0.eventLoop.makeSucceededVoidFuture()
  142. }
  143. XCTAssertThrowsError(try stream.wait()) { error in
  144. XCTAssert((error as? ConnectionPoolError).isShutdown)
  145. }
  146. }
  147. func testMakeStreamWhenWaiterQueueIsFull() {
  148. let maxWaiters = 5
  149. let pool = self.makePool(waiters: maxWaiters) {
  150. self.noChannelExpected($0, $1)
  151. }
  152. let waiting = (0 ..< maxWaiters).map { _ in
  153. return pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  154. $0.eventLoop.makeSucceededVoidFuture()
  155. }
  156. }
  157. let tooManyWaiters = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  158. $0.eventLoop.makeSucceededVoidFuture()
  159. }
  160. XCTAssertThrowsError(try tooManyWaiters.wait()) { error in
  161. XCTAssert((error as? ConnectionPoolError).isTooManyWaiters)
  162. }
  163. XCTAssertNoThrow(try pool.shutdown().wait())
  164. // All 'waiting' futures will be failed by the shutdown promise.
  165. for waiter in waiting {
  166. XCTAssertThrowsError(try waiter.wait()) { error in
  167. XCTAssert((error as? ConnectionPoolError).isShutdown)
  168. }
  169. }
  170. }
  171. func testWaiterTimingOut() {
  172. let pool = self.makePool {
  173. self.noChannelExpected($0, $1)
  174. }
  175. let waiter = pool.makeStream(deadline: .uptimeNanoseconds(10), logger: self.logger.wrapped) {
  176. $0.eventLoop.makeSucceededVoidFuture()
  177. }
  178. XCTAssertEqual(pool.sync.waiters, 1)
  179. self.eventLoop.advanceTime(to: .uptimeNanoseconds(10))
  180. XCTAssertThrowsError(try waiter.wait()) { error in
  181. XCTAssert((error as? ConnectionPoolError).isDeadlineExceeded)
  182. }
  183. XCTAssertEqual(pool.sync.waiters, 0)
  184. }
  185. func testWaiterTimingOutInPast() {
  186. let pool = self.makePool {
  187. self.noChannelExpected($0, $1)
  188. }
  189. self.eventLoop.advanceTime(to: .uptimeNanoseconds(10))
  190. let waiter = pool.makeStream(deadline: .uptimeNanoseconds(5), logger: self.logger.wrapped) {
  191. $0.eventLoop.makeSucceededVoidFuture()
  192. }
  193. XCTAssertEqual(pool.sync.waiters, 1)
  194. self.eventLoop.run()
  195. XCTAssertThrowsError(try waiter.wait()) { error in
  196. XCTAssert((error as? ConnectionPoolError).isDeadlineExceeded)
  197. }
  198. XCTAssertEqual(pool.sync.waiters, 0)
  199. }
  200. func testMakeStreamTriggersChannelCreation() {
  201. let (pool, controller) = self.setUpPoolAndController()
  202. pool.initialize(connections: 1)
  203. XCTAssertEqual(pool.sync.connections, 1)
  204. // No channels yet.
  205. XCTAssertEqual(controller.count, 0)
  206. let waiter = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  207. $0.eventLoop.makeSucceededVoidFuture()
  208. }
  209. // Start creating the channel.
  210. self.eventLoop.run()
  211. // We should have been asked for a channel now.
  212. XCTAssertEqual(controller.count, 1)
  213. // The connection isn't ready yet though, so no streams available.
  214. XCTAssertEqual(pool.sync.availableStreams, 0)
  215. // Make the connection 'ready'.
  216. controller.connectChannel(atIndex: 0)
  217. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: 10)
  218. // We have a multiplexer and a 'ready' connection.
  219. XCTAssertEqual(pool.sync.reservedStreams, 1)
  220. XCTAssertEqual(pool.sync.availableStreams, 9)
  221. XCTAssertEqual(pool.sync.waiters, 0)
  222. // Run the loop to create the stream, we need to fire the event too.
  223. self.eventLoop.run()
  224. XCTAssertNoThrow(try waiter.wait())
  225. controller.openStreamInChannel(atIndex: 0)
  226. // Now close the stream.
  227. controller.closeStreamInChannel(atIndex: 0)
  228. XCTAssertEqual(pool.sync.reservedStreams, 0)
  229. XCTAssertEqual(pool.sync.availableStreams, 10)
  230. }
  231. func testMakeStreamWhenConnectionIsAlreadyAvailable() {
  232. let (pool, controller) = self.setUpPoolAndController()
  233. pool.initialize(connections: 1)
  234. let waiter = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  235. $0.eventLoop.makeSucceededVoidFuture()
  236. }
  237. // Start creating the channel.
  238. self.eventLoop.run()
  239. XCTAssertEqual(controller.count, 1)
  240. // Fire up the connection.
  241. controller.connectChannel(atIndex: 0)
  242. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: 10)
  243. // Run the loop to create the stream, we need to fire the stream creation event too.
  244. self.eventLoop.run()
  245. XCTAssertNoThrow(try waiter.wait())
  246. controller.openStreamInChannel(atIndex: 0)
  247. // Now we can create another stream, but as there's already an available stream on an active
  248. // connection we won't have to wait.
  249. XCTAssertEqual(pool.sync.waiters, 0)
  250. XCTAssertEqual(pool.sync.reservedStreams, 1)
  251. let notWaiting = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  252. $0.eventLoop.makeSucceededVoidFuture()
  253. }
  254. // Still no waiters.
  255. XCTAssertEqual(pool.sync.waiters, 0)
  256. XCTAssertEqual(pool.sync.reservedStreams, 2)
  257. // Run the loop to create the stream, we need to fire the stream creation event too.
  258. self.eventLoop.run()
  259. XCTAssertNoThrow(try notWaiting.wait())
  260. controller.openStreamInChannel(atIndex: 0)
  261. }
  262. func testMakeMoreWaitersThanConnectionCanHandle() {
  263. var returnedStreams: [Int] = []
  264. let (pool, controller) = self.setUpPoolAndController(onReservationReturned: {
  265. returnedStreams.append($0)
  266. })
  267. pool.initialize(connections: 1)
  268. // Enqueue twice as many waiters as the connection will be able to handle.
  269. let maxConcurrentStreams = 10
  270. let waiters = (0 ..< maxConcurrentStreams * 2).map { _ in
  271. return pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  272. $0.eventLoop.makeSucceededVoidFuture()
  273. }
  274. }
  275. XCTAssertEqual(pool.sync.waiters, 2 * maxConcurrentStreams)
  276. // Fire up the connection.
  277. self.eventLoop.run()
  278. controller.connectChannel(atIndex: 0)
  279. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: maxConcurrentStreams)
  280. // We should have assigned a bunch of streams to waiters now.
  281. XCTAssertEqual(pool.sync.waiters, maxConcurrentStreams)
  282. XCTAssertEqual(pool.sync.reservedStreams, maxConcurrentStreams)
  283. XCTAssertEqual(pool.sync.availableStreams, 0)
  284. // Do the stream creation and make sure the first batch are succeeded.
  285. self.eventLoop.run()
  286. let firstBatch = waiters.prefix(maxConcurrentStreams)
  287. var others = waiters.dropFirst(maxConcurrentStreams)
  288. for waiter in firstBatch {
  289. XCTAssertNoThrow(try waiter.wait())
  290. controller.openStreamInChannel(atIndex: 0)
  291. }
  292. // Close a stream.
  293. controller.closeStreamInChannel(atIndex: 0)
  294. XCTAssertEqual(returnedStreams, [1])
  295. // We have another stream so a waiter should be succeeded.
  296. XCTAssertEqual(pool.sync.waiters, maxConcurrentStreams - 1)
  297. self.eventLoop.run()
  298. XCTAssertNoThrow(try others.popFirst()?.wait())
  299. // Shutdown the pool: the remaining waiters should be failed.
  300. let shutdown = pool.shutdown()
  301. self.eventLoop.run()
  302. XCTAssertNoThrow(try shutdown.wait())
  303. for waiter in others {
  304. XCTAssertThrowsError(try waiter.wait()) { error in
  305. XCTAssert((error as? ConnectionPoolError).isShutdown)
  306. }
  307. }
  308. }
  309. func testDropConnectionWithOutstandingReservations() {
  310. var streamsReturned: [Int] = []
  311. let (pool, controller) = self.setUpPoolAndController(
  312. onReservationReturned: { streamsReturned.append($0) }
  313. )
  314. pool.initialize(connections: 1)
  315. let waiter = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  316. $0.eventLoop.makeSucceededVoidFuture()
  317. }
  318. // Start creating the channel.
  319. self.eventLoop.run()
  320. XCTAssertEqual(controller.count, 1)
  321. // Fire up the connection.
  322. controller.connectChannel(atIndex: 0)
  323. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: 10)
  324. // Run the loop to create the stream, we need to fire the stream creation event too.
  325. self.eventLoop.run()
  326. XCTAssertNoThrow(try waiter.wait())
  327. controller.openStreamInChannel(atIndex: 0)
  328. // Create a handful of streams.
  329. XCTAssertEqual(pool.sync.availableStreams, 9)
  330. for _ in 0 ..< 5 {
  331. let notWaiting = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  332. $0.eventLoop.makeSucceededVoidFuture()
  333. }
  334. self.eventLoop.run()
  335. XCTAssertNoThrow(try notWaiting.wait())
  336. controller.openStreamInChannel(atIndex: 0)
  337. }
  338. XCTAssertEqual(pool.sync.availableStreams, 4)
  339. XCTAssertEqual(pool.sync.reservedStreams, 6)
  340. // Blast the connection away. We'll be notified about dropped reservations.
  341. XCTAssertEqual(streamsReturned, [])
  342. controller.throwError(ChannelError.ioOnClosedChannel, inChannelAtIndex: 0)
  343. controller.fireChannelInactiveForChannel(atIndex: 0)
  344. XCTAssertEqual(streamsReturned, [6])
  345. XCTAssertEqual(pool.sync.availableStreams, 0)
  346. XCTAssertEqual(pool.sync.reservedStreams, 0)
  347. }
  348. func testDropConnectionWithOutstandingReservationsAndWaiters() {
  349. var streamsReturned: [Int] = []
  350. let (pool, controller) = self.setUpPoolAndController(
  351. onReservationReturned: { streamsReturned.append($0) }
  352. )
  353. pool.initialize(connections: 1)
  354. // Reserve a bunch of streams.
  355. let waiters = (0 ..< 10).map { _ in
  356. return pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  357. $0.eventLoop.makeSucceededVoidFuture()
  358. }
  359. }
  360. // Connect and setup all the streams.
  361. self.eventLoop.run()
  362. controller.connectChannel(atIndex: 0)
  363. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: 10)
  364. self.eventLoop.run()
  365. for waiter in waiters {
  366. XCTAssertNoThrow(try waiter.wait())
  367. controller.openStreamInChannel(atIndex: 0)
  368. }
  369. // All streams should be reserved.
  370. XCTAssertEqual(pool.sync.availableStreams, 0)
  371. XCTAssertEqual(pool.sync.reservedStreams, 10)
  372. // Add a waiter.
  373. XCTAssertEqual(pool.sync.waiters, 0)
  374. let waiter = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  375. $0.eventLoop.makeSucceededVoidFuture()
  376. }
  377. XCTAssertEqual(pool.sync.waiters, 1)
  378. // Now bork the connection. We'll be notified about the 10 dropped reservation but not the one
  379. // waiter .
  380. XCTAssertEqual(streamsReturned, [])
  381. controller.throwError(ChannelError.ioOnClosedChannel, inChannelAtIndex: 0)
  382. controller.fireChannelInactiveForChannel(atIndex: 0)
  383. XCTAssertEqual(streamsReturned, [10])
  384. // The connection dropped, let the reconnect kick in.
  385. self.eventLoop.run()
  386. XCTAssertEqual(controller.count, 2)
  387. controller.connectChannel(atIndex: 1)
  388. controller.sendSettingsToChannel(atIndex: 1, maxConcurrentStreams: 10)
  389. self.eventLoop.run()
  390. XCTAssertNoThrow(try waiter.wait())
  391. controller.openStreamInChannel(atIndex: 1)
  392. controller.closeStreamInChannel(atIndex: 1)
  393. XCTAssertEqual(streamsReturned, [10, 1])
  394. XCTAssertEqual(pool.sync.availableStreams, 10)
  395. XCTAssertEqual(pool.sync.reservedStreams, 0)
  396. }
  397. func testDeadlineExceededInSameTickAsSucceedingWaiters() {
  398. // deadline must be exceeded just as servicing waiter is done
  399. // - setup waiter with deadline x
  400. // - start connecting
  401. // - set time to x
  402. // - finish connecting
  403. let (pool, controller) = self.setUpPoolAndController(now: {
  404. return NIODeadline.uptimeNanoseconds(12)
  405. })
  406. pool.initialize(connections: 1)
  407. let waiter1 = pool.makeStream(deadline: .uptimeNanoseconds(10), logger: self.logger.wrapped) {
  408. $0.eventLoop.makeSucceededVoidFuture()
  409. }
  410. let waiter2 = pool.makeStream(deadline: .uptimeNanoseconds(15), logger: self.logger.wrapped) {
  411. $0.eventLoop.makeSucceededVoidFuture()
  412. }
  413. // Start creating the channel.
  414. self.eventLoop.run()
  415. XCTAssertEqual(controller.count, 1)
  416. // Fire up the connection.
  417. controller.connectChannel(atIndex: 0)
  418. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: 10)
  419. // The deadline for the first waiter is already after 'now', so it'll fail with deadline
  420. // exceeded.
  421. self.eventLoop.run()
  422. // We need to advance the time to fire the timeout to fail the waiter.
  423. self.eventLoop.advanceTime(to: .uptimeNanoseconds(10))
  424. XCTAssertThrowsError(try waiter1.wait()) { error in
  425. XCTAssert((error as? ConnectionPoolError).isDeadlineExceeded)
  426. }
  427. self.eventLoop.run()
  428. XCTAssertNoThrow(try waiter2.wait())
  429. controller.openStreamInChannel(atIndex: 0)
  430. XCTAssertEqual(pool.sync.waiters, 0)
  431. XCTAssertEqual(pool.sync.reservedStreams, 1)
  432. XCTAssertEqual(pool.sync.availableStreams, 9)
  433. controller.closeStreamInChannel(atIndex: 0)
  434. XCTAssertEqual(pool.sync.waiters, 0)
  435. XCTAssertEqual(pool.sync.reservedStreams, 0)
  436. XCTAssertEqual(pool.sync.availableStreams, 10)
  437. }
  438. func testConnectionsAreBroughtUpAtAppropriateTimes() {
  439. let (pool, controller) = self.setUpPoolAndController(reservationLoadThreshold: 0.2)
  440. // We'll allow 3 connections and configure max concurrent streams to 10. With our reservation
  441. // threshold we'll bring up a new connection after enqueueing the 1st, 2nd and 4th waiters.
  442. pool.initialize(connections: 3)
  443. let maxConcurrentStreams = 10
  444. // No demand so all three connections are idle.
  445. XCTAssertEqual(pool.sync.idleConnections, 3)
  446. let w1 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  447. $0.eventLoop.makeSucceededVoidFuture()
  448. }
  449. // demand=1, available=0, load=infinite, one connection should be non-idle
  450. XCTAssertEqual(pool.sync.idleConnections, 2)
  451. // Connect the first channel and write the first settings frame; this allows us to lower the
  452. // default max concurrent streams value (from 100).
  453. self.eventLoop.run()
  454. controller.connectChannel(atIndex: 0)
  455. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: maxConcurrentStreams)
  456. self.eventLoop.run()
  457. XCTAssertNoThrow(try w1.wait())
  458. controller.openStreamInChannel(atIndex: 0)
  459. let w2 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  460. $0.eventLoop.makeSucceededVoidFuture()
  461. }
  462. self.eventLoop.run()
  463. XCTAssertNoThrow(try w2.wait())
  464. controller.openStreamInChannel(atIndex: 0)
  465. // demand=2, available=10, load=0.2; only one idle connection now.
  466. XCTAssertEqual(pool.sync.idleConnections, 1)
  467. // Add more demand before the second connection comes up.
  468. let w3 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  469. $0.eventLoop.makeSucceededVoidFuture()
  470. }
  471. // demand=3, available=20, load=0.15; still one idle connection.
  472. XCTAssertEqual(pool.sync.idleConnections, 1)
  473. // Connection the next channel
  474. self.eventLoop.run()
  475. controller.connectChannel(atIndex: 1)
  476. controller.sendSettingsToChannel(atIndex: 1, maxConcurrentStreams: maxConcurrentStreams)
  477. XCTAssertNoThrow(try w3.wait())
  478. controller.openStreamInChannel(atIndex: 1)
  479. }
  480. func testQuiescingConnectionIsReplaced() {
  481. var reservationsReturned: [Int] = []
  482. let (pool, controller) = self.setUpPoolAndController(onReservationReturned: {
  483. reservationsReturned.append($0)
  484. })
  485. pool.initialize(connections: 1)
  486. XCTAssertEqual(pool.sync.connections, 1)
  487. let w1 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  488. $0.eventLoop.makeSucceededVoidFuture()
  489. }
  490. // Start creating the channel.
  491. self.eventLoop.run()
  492. // Make the connection 'ready'.
  493. controller.connectChannel(atIndex: 0)
  494. controller.sendSettingsToChannel(atIndex: 0)
  495. // Run the loop to create the stream.
  496. self.eventLoop.run()
  497. XCTAssertNoThrow(try w1.wait())
  498. controller.openStreamInChannel(atIndex: 0)
  499. // One stream reserved by 'w1' on the only connection in the pool (which isn't idle).
  500. XCTAssertEqual(pool.sync.reservedStreams, 1)
  501. XCTAssertEqual(pool.sync.connections, 1)
  502. XCTAssertEqual(pool.sync.idleConnections, 0)
  503. // Quiesce the connection. It should be punted from the pool and any active RPCs allowed to run
  504. // their course. A new (idle) connection should replace it in the pool.
  505. controller.sendGoAwayToChannel(atIndex: 0)
  506. // The quiescing connection had 1 stream reserved, it's now returned to the outer pool and we
  507. // have a new idle connection in place of the old one.
  508. XCTAssertEqual(reservationsReturned, [1])
  509. // The inner pool still knows about the reserved stream.
  510. XCTAssertEqual(pool.sync.reservedStreams, 1)
  511. XCTAssertEqual(pool.sync.availableStreams, 0)
  512. XCTAssertEqual(pool.sync.idleConnections, 1)
  513. // Ask for another stream: this will be on the new idle connection.
  514. let w2 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  515. $0.eventLoop.makeSucceededVoidFuture()
  516. }
  517. self.eventLoop.run()
  518. XCTAssertEqual(controller.count, 2)
  519. // Make the connection 'ready'.
  520. controller.connectChannel(atIndex: 1)
  521. controller.sendSettingsToChannel(atIndex: 1)
  522. self.eventLoop.run()
  523. XCTAssertNoThrow(try w2.wait())
  524. controller.openStreamInChannel(atIndex: 1)
  525. // The stream on the quiescing connection is still reserved.
  526. XCTAssertEqual(pool.sync.reservedStreams, 2)
  527. XCTAssertEqual(pool.sync.availableStreams, 99)
  528. // Return a stream for the _quiescing_ connection: nothing should change in the pool.
  529. controller.closeStreamInChannel(atIndex: 0)
  530. XCTAssertEqual(pool.sync.reservedStreams, 1)
  531. XCTAssertEqual(pool.sync.availableStreams, 99)
  532. // Return a stream for the new connection.
  533. controller.closeStreamInChannel(atIndex: 1)
  534. XCTAssertEqual(reservationsReturned, [1, 1])
  535. XCTAssertEqual(pool.sync.reservedStreams, 0)
  536. XCTAssertEqual(pool.sync.availableStreams, 100)
  537. }
  538. func testBackoffIsUsedForReconnections() {
  539. // Fix backoff to always be 1 second.
  540. let backoff = ConnectionBackoff(
  541. initialBackoff: 1.0,
  542. maximumBackoff: 1.0,
  543. multiplier: 1.0,
  544. jitter: 0.0
  545. )
  546. let (pool, controller) = self.setUpPoolAndController(connectionBackoff: backoff)
  547. pool.initialize(connections: 1)
  548. XCTAssertEqual(pool.sync.connections, 1)
  549. let w1 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  550. $0.eventLoop.makeSucceededVoidFuture()
  551. }
  552. // Start creating the channel.
  553. self.eventLoop.run()
  554. // Make the connection 'ready'.
  555. controller.connectChannel(atIndex: 0)
  556. controller.sendSettingsToChannel(atIndex: 0)
  557. self.eventLoop.run()
  558. XCTAssertNoThrow(try w1.wait())
  559. controller.openStreamInChannel(atIndex: 0)
  560. // Close the connection. It should hit the transient failure state.
  561. controller.fireChannelInactiveForChannel(atIndex: 0)
  562. // Now nothing is available in the pool.
  563. XCTAssertEqual(pool.sync.waiters, 0)
  564. XCTAssertEqual(pool.sync.availableStreams, 0)
  565. XCTAssertEqual(pool.sync.reservedStreams, 0)
  566. XCTAssertEqual(pool.sync.idleConnections, 0)
  567. // Enqueue two waiters. One to time out before the reconnect happens.
  568. let w2 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  569. $0.eventLoop.makeSucceededVoidFuture()
  570. }
  571. let w3 = pool.makeStream(
  572. deadline: .uptimeNanoseconds(UInt64(TimeAmount.milliseconds(500).nanoseconds)),
  573. logger: self.logger.wrapped
  574. ) {
  575. $0.eventLoop.makeSucceededVoidFuture()
  576. }
  577. XCTAssertEqual(pool.sync.waiters, 2)
  578. // Time out w3.
  579. self.eventLoop.advanceTime(by: .milliseconds(500))
  580. XCTAssertThrowsError(try w3.wait())
  581. XCTAssertEqual(pool.sync.waiters, 1)
  582. // Wait a little more for the backoff to pass. The controller should now have a second channel.
  583. self.eventLoop.advanceTime(by: .milliseconds(500))
  584. XCTAssertEqual(controller.count, 2)
  585. // Start up the next channel.
  586. controller.connectChannel(atIndex: 1)
  587. controller.sendSettingsToChannel(atIndex: 1)
  588. self.eventLoop.run()
  589. XCTAssertNoThrow(try w2.wait())
  590. controller.openStreamInChannel(atIndex: 1)
  591. }
  592. func testFailedWaiterWithError() throws {
  593. // We want to check a few things in this test:
  594. //
  595. // 1. When an active channel throws an error that any waiter in the connection pool which has
  596. // its deadline exceeded or any waiter which exceeds the waiter limit fails with an error
  597. // which includes the underlying channel error.
  598. // 2. When a reconnect happens and the pool is just busy, no underlying error is passed through
  599. // to failing waiters.
  600. // Fix backoff to always be 1 second. This is necessary to figure out timings later on when
  601. // we try to establish a new connection.
  602. let backoff = ConnectionBackoff(
  603. initialBackoff: 1.0,
  604. maximumBackoff: 1.0,
  605. multiplier: 1.0,
  606. jitter: 0.0
  607. )
  608. let (pool, controller) = self.setUpPoolAndController(waiters: 10, connectionBackoff: backoff)
  609. pool.initialize(connections: 1)
  610. // First we'll create two streams which will fail for different reasons.
  611. // - w1 will fail because of a timeout (no channel came up before the waiters own deadline
  612. // passed but no connection has previously failed)
  613. // - w2 will fail because of a timeout but after the underlying channel has failed to connect so
  614. // should have that additional failure information.
  615. let w1 = pool.makeStream(deadline: .uptimeNanoseconds(10), logger: self.logger.wrapped) {
  616. $0.eventLoop.makeSucceededVoidFuture()
  617. }
  618. let w2 = pool.makeStream(deadline: .uptimeNanoseconds(20), logger: self.logger.wrapped) {
  619. $0.eventLoop.makeSucceededVoidFuture()
  620. }
  621. // Start creating the channel.
  622. self.eventLoop.run()
  623. XCTAssertEqual(controller.count, 1)
  624. // Fire up the connection.
  625. controller.connectChannel(atIndex: 0)
  626. // Advance time to fail the w1.
  627. self.eventLoop.advanceTime(to: .uptimeNanoseconds(10))
  628. XCTAssertThrowsError(try w1.wait()) { error in
  629. switch error as? ConnectionPoolError {
  630. case .some(.deadlineExceeded(.none)):
  631. // Deadline exceeded but no underlying error, as expected.
  632. ()
  633. default:
  634. XCTFail("Expected ConnectionPoolError.deadlineExceeded(.none) but got \(error)")
  635. }
  636. }
  637. // Now fail the connection and timeout w2.
  638. struct DummyError: Error {}
  639. controller.throwError(DummyError(), inChannelAtIndex: 0)
  640. controller.fireChannelInactiveForChannel(atIndex: 0)
  641. self.eventLoop.advanceTime(to: .uptimeNanoseconds(20))
  642. XCTAssertThrowsError(try w2.wait()) { error in
  643. switch error as? ConnectionPoolError {
  644. case let .some(.deadlineExceeded(.some(wrappedError))):
  645. // Deadline exceeded and we have the underlying error.
  646. XCTAssert(wrappedError is DummyError)
  647. default:
  648. XCTFail("Expected ConnectionPoolError.deadlineExceeded(.some) but got \(error)")
  649. }
  650. }
  651. // For the next part of the test we want to validate that when a new channel is created after
  652. // the backoff period passes that no additional errors are attached when the pool is just busy
  653. // but otherwise operational.
  654. //
  655. // To do this we'll create a bunch of waiters. These will be succeeded when the new connection
  656. // comes up and, importantly, use up all available streams on that connection.
  657. //
  658. // We'll then enqueue enough waiters to fill the waiter queue. We'll then validate that one more
  659. // waiter trips over the queue limit but does not include the connection error we saw earlier.
  660. // We'll then timeout the waiters in the queue and validate the same thing.
  661. // These streams should succeed when the new connection is up. We'll limit the connection to 10
  662. // streams when we bring it up.
  663. let streams = (0 ..< 10).map { _ in
  664. pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  665. $0.eventLoop.makeSucceededVoidFuture()
  666. }
  667. }
  668. // The connection is backing off; advance time to create another channel.
  669. XCTAssertEqual(controller.count, 1)
  670. self.eventLoop.advanceTime(by: .seconds(1))
  671. XCTAssertEqual(controller.count, 2)
  672. controller.connectChannel(atIndex: 1)
  673. controller.sendSettingsToChannel(atIndex: 1, maxConcurrentStreams: 10)
  674. self.eventLoop.run()
  675. // Make sure the streams are succeeded.
  676. for stream in streams {
  677. XCTAssertNoThrow(try stream.wait())
  678. controller.openStreamInChannel(atIndex: 1)
  679. }
  680. // All streams should be reserved.
  681. XCTAssertEqual(pool.sync.availableStreams, 0)
  682. XCTAssertEqual(pool.sync.reservedStreams, 10)
  683. XCTAssertEqual(pool.sync.waiters, 0)
  684. // We configured the pool to allow for 10 waiters, so let's enqueue that many which will time
  685. // out at a known point in time.
  686. let now = NIODeadline.now()
  687. self.eventLoop.advanceTime(to: now)
  688. let waiters = (0 ..< 10).map { _ in
  689. pool.makeStream(deadline: now + .seconds(1), logger: self.logger.wrapped) {
  690. $0.eventLoop.makeSucceededVoidFuture()
  691. }
  692. }
  693. // This is one waiter more than is allowed so it should hit too-many-waiters. We don't expect
  694. // an inner error though, the connection is just busy.
  695. let tooManyWaiters = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  696. $0.eventLoop.makeSucceededVoidFuture()
  697. }
  698. XCTAssertThrowsError(try tooManyWaiters.wait()) { error in
  699. switch error as? ConnectionPoolError {
  700. case .some(.tooManyWaiters(.none)):
  701. ()
  702. default:
  703. XCTFail("Expected ConnectionPoolError.tooManyWaiters(.none) but got \(error)")
  704. }
  705. }
  706. // Finally, timeout the remaining waiters. Again, no inner error, the connection is just busy.
  707. self.eventLoop.advanceTime(by: .seconds(1))
  708. for waiter in waiters {
  709. XCTAssertThrowsError(try waiter.wait()) { error in
  710. switch error as? ConnectionPoolError {
  711. case .some(.deadlineExceeded(.none)):
  712. ()
  713. default:
  714. XCTFail("Expected ConnectionPoolError.deadlineExceeded(.none) but got \(error)")
  715. }
  716. }
  717. }
  718. }
  719. func testWaiterStoresItsScheduledTask() throws {
  720. let deadline = NIODeadline.uptimeNanoseconds(42)
  721. let promise = self.eventLoop.makePromise(of: Channel.self)
  722. let waiter = ConnectionPool.Waiter(deadline: deadline, promise: promise) {
  723. return $0.eventLoop.makeSucceededVoidFuture()
  724. }
  725. XCTAssertNil(waiter._scheduledTimeout)
  726. waiter.scheduleTimeout(on: self.eventLoop) {
  727. waiter.fail(ConnectionPoolError.deadlineExceeded(connectionError: nil))
  728. }
  729. XCTAssertNotNil(waiter._scheduledTimeout)
  730. self.eventLoop.advanceTime(to: deadline)
  731. XCTAssertThrowsError(try promise.futureResult.wait())
  732. XCTAssertNil(waiter._scheduledTimeout)
  733. }
  734. func testConnectionPoolDelegate() throws {
  735. let recorder = EventRecordingConnectionPoolDelegate()
  736. let (pool, controller) = self.setUpPoolAndController(delegate: recorder)
  737. pool.initialize(connections: 2)
  738. func assertConnectionAdded(
  739. _ event: EventRecordingConnectionPoolDelegate.Event?
  740. ) throws -> GRPCConnectionID {
  741. let unwrappedEvent = try XCTUnwrap(event)
  742. switch unwrappedEvent {
  743. case let .connectionAdded(id):
  744. return id
  745. default:
  746. throw EventRecordingConnectionPoolDelegate.UnexpectedEvent(unwrappedEvent)
  747. }
  748. }
  749. let connID1 = try assertConnectionAdded(recorder.popFirst())
  750. let connID2 = try assertConnectionAdded(recorder.popFirst())
  751. let waiter = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  752. $0.eventLoop.makeSucceededVoidFuture()
  753. }
  754. // Start creating the channel.
  755. self.eventLoop.run()
  756. let startedConnecting = recorder.popFirst()
  757. let firstConn: GRPCConnectionID
  758. let secondConn: GRPCConnectionID
  759. if startedConnecting == .startedConnecting(connID1) {
  760. firstConn = connID1
  761. secondConn = connID2
  762. } else if startedConnecting == .startedConnecting(connID2) {
  763. firstConn = connID2
  764. secondConn = connID1
  765. } else {
  766. return XCTFail("Unexpected event")
  767. }
  768. // Connect the connection.
  769. self.eventLoop.run()
  770. controller.connectChannel(atIndex: 0)
  771. controller.sendSettingsToChannel(atIndex: 0, maxConcurrentStreams: 10)
  772. XCTAssertEqual(recorder.popFirst(), .connectSucceeded(firstConn, 10))
  773. // Open a stream for the waiter.
  774. controller.openStreamInChannel(atIndex: 0)
  775. XCTAssertEqual(recorder.popFirst(), .connectionUtilizationChanged(firstConn, 1, 10))
  776. self.eventLoop.run()
  777. XCTAssertNoThrow(try waiter.wait())
  778. // Okay, more utilization!
  779. for n in 2 ... 8 {
  780. let w = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  781. $0.eventLoop.makeSucceededVoidFuture()
  782. }
  783. controller.openStreamInChannel(atIndex: 0)
  784. XCTAssertEqual(recorder.popFirst(), .connectionUtilizationChanged(firstConn, n, 10))
  785. self.eventLoop.run()
  786. XCTAssertNoThrow(try w.wait())
  787. }
  788. // The utilisation threshold before bringing up a new connection is 0.9; we have 8 open streams
  789. // (out of 10) now so opening the next should trigger a connect on the other connection.
  790. let w9 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  791. $0.eventLoop.makeSucceededVoidFuture()
  792. }
  793. XCTAssertEqual(recorder.popFirst(), .startedConnecting(secondConn))
  794. // Deal with the 9th stream.
  795. controller.openStreamInChannel(atIndex: 0)
  796. XCTAssertEqual(recorder.popFirst(), .connectionUtilizationChanged(firstConn, 9, 10))
  797. self.eventLoop.run()
  798. XCTAssertNoThrow(try w9.wait())
  799. // Bring up the next connection.
  800. controller.connectChannel(atIndex: 1)
  801. controller.sendSettingsToChannel(atIndex: 1, maxConcurrentStreams: 10)
  802. XCTAssertEqual(recorder.popFirst(), .connectSucceeded(secondConn, 10))
  803. // The next stream should be on the new connection.
  804. let w10 = pool.makeStream(deadline: .distantFuture, logger: self.logger.wrapped) {
  805. $0.eventLoop.makeSucceededVoidFuture()
  806. }
  807. // Deal with the 10th stream.
  808. controller.openStreamInChannel(atIndex: 1)
  809. XCTAssertEqual(recorder.popFirst(), .connectionUtilizationChanged(secondConn, 1, 10))
  810. self.eventLoop.run()
  811. XCTAssertNoThrow(try w10.wait())
  812. // Close the streams.
  813. for i in 1 ... 9 {
  814. controller.closeStreamInChannel(atIndex: 0)
  815. XCTAssertEqual(recorder.popFirst(), .connectionUtilizationChanged(firstConn, 9 - i, 10))
  816. }
  817. controller.closeStreamInChannel(atIndex: 1)
  818. XCTAssertEqual(recorder.popFirst(), .connectionUtilizationChanged(secondConn, 0, 10))
  819. // Close the connections.
  820. controller.fireChannelInactiveForChannel(atIndex: 0)
  821. XCTAssertEqual(recorder.popFirst(), .connectionClosed(firstConn))
  822. controller.fireChannelInactiveForChannel(atIndex: 1)
  823. XCTAssertEqual(recorder.popFirst(), .connectionClosed(secondConn))
  824. // All conns are already closed.
  825. let shutdownFuture = pool.shutdown()
  826. self.eventLoop.run()
  827. XCTAssertNoThrow(try shutdownFuture.wait())
  828. // Two connections must be removed.
  829. for _ in 0 ..< 2 {
  830. if let event = recorder.popFirst() {
  831. let id = event.id
  832. XCTAssertEqual(event, .connectionRemoved(id))
  833. } else {
  834. XCTFail("Expected .connectionRemoved")
  835. }
  836. }
  837. }
  838. }
  839. extension ConnectionPool {
  840. // For backwards compatibility, to avoid large diffs in these tests.
  841. fileprivate func shutdown() -> EventLoopFuture<Void> {
  842. return self.shutdown(mode: .forceful)
  843. }
  844. }
  845. // MARK: - Helpers
  846. struct ChannelAndState {
  847. let channel: EmbeddedChannel
  848. let streamDelegate: NIOHTTP2StreamDelegate
  849. var isActive: Bool
  850. }
  851. internal final class ChannelController {
  852. private var channels: [ChannelAndState] = []
  853. internal var count: Int {
  854. return self.channels.count
  855. }
  856. internal func finish() {
  857. while let state = self.channels.popLast() {
  858. if state.isActive {
  859. _ = try? state.channel.finish()
  860. }
  861. }
  862. }
  863. private func isValidIndex(
  864. _ index: Int,
  865. file: StaticString = #filePath,
  866. line: UInt = #line
  867. ) -> Bool {
  868. let isValid = self.channels.indices.contains(index)
  869. XCTAssertTrue(isValid, "Invalid connection index '\(index)'", file: file, line: line)
  870. return isValid
  871. }
  872. internal func connectChannel(
  873. atIndex index: Int,
  874. file: StaticString = #filePath,
  875. line: UInt = #line
  876. ) {
  877. guard self.isValidIndex(index, file: file, line: line) else { return }
  878. self.channels[index].isActive = true
  879. XCTAssertNoThrow(
  880. try self.channels[index].channel.connect(to: .init(unixDomainSocketPath: "/")),
  881. file: file,
  882. line: line
  883. )
  884. }
  885. internal func fireChannelInactiveForChannel(
  886. atIndex index: Int,
  887. file: StaticString = #filePath,
  888. line: UInt = #line
  889. ) {
  890. guard self.isValidIndex(index, file: file, line: line) else { return }
  891. self.channels[index].channel.pipeline.fireChannelInactive()
  892. self.channels[index].isActive = false
  893. }
  894. internal func throwError(
  895. _ error: Error,
  896. inChannelAtIndex index: Int,
  897. file: StaticString = #filePath,
  898. line: UInt = #line
  899. ) {
  900. guard self.isValidIndex(index, file: file, line: line) else { return }
  901. self.channels[index].channel.pipeline.fireErrorCaught(error)
  902. }
  903. internal func sendSettingsToChannel(
  904. atIndex index: Int,
  905. maxConcurrentStreams: Int = 100,
  906. file: StaticString = #filePath,
  907. line: UInt = #line
  908. ) {
  909. guard self.isValidIndex(index, file: file, line: line) else { return }
  910. let settings = [HTTP2Setting(parameter: .maxConcurrentStreams, value: maxConcurrentStreams)]
  911. let settingsFrame = HTTP2Frame(streamID: .rootStream, payload: .settings(.settings(settings)))
  912. XCTAssertNoThrow(
  913. try self.channels[index].channel.writeInbound(settingsFrame.encode()),
  914. file: file,
  915. line: line
  916. )
  917. }
  918. internal func sendGoAwayToChannel(
  919. atIndex index: Int,
  920. file: StaticString = #filePath,
  921. line: UInt = #line
  922. ) {
  923. guard self.isValidIndex(index, file: file, line: line) else { return }
  924. let goAwayFrame = HTTP2Frame(
  925. streamID: .rootStream,
  926. payload: .goAway(lastStreamID: .maxID, errorCode: .noError, opaqueData: nil)
  927. )
  928. XCTAssertNoThrow(
  929. try self.channels[index].channel.writeInbound(goAwayFrame.encode()),
  930. file: file,
  931. line: line
  932. )
  933. }
  934. internal func openStreamInChannel(
  935. atIndex index: Int,
  936. file: StaticString = #filePath,
  937. line: UInt = #line
  938. ) {
  939. guard self.isValidIndex(index, file: file, line: line) else { return }
  940. // The details don't matter here.
  941. let channel = self.channels[index]
  942. channel.streamDelegate.streamCreated(.rootStream, channel: channel.channel)
  943. }
  944. internal func closeStreamInChannel(
  945. atIndex index: Int,
  946. file: StaticString = #filePath,
  947. line: UInt = #line
  948. ) {
  949. guard self.isValidIndex(index, file: file, line: line) else { return }
  950. // The details don't matter here.
  951. let channel = self.channels[index]
  952. channel.streamDelegate.streamClosed(.rootStream, channel: channel.channel)
  953. }
  954. }
  955. extension ChannelController: ConnectionManagerChannelProvider {
  956. internal func makeChannel(
  957. managedBy connectionManager: ConnectionManager,
  958. onEventLoop eventLoop: EventLoop,
  959. connectTimeout: TimeAmount?,
  960. logger: Logger
  961. ) -> EventLoopFuture<Channel> {
  962. let channel = EmbeddedChannel(loop: eventLoop as! EmbeddedEventLoop)
  963. let idleHandler = GRPCIdleHandler(
  964. connectionManager: connectionManager,
  965. idleTimeout: .minutes(5),
  966. keepalive: ClientConnectionKeepalive(),
  967. logger: logger
  968. )
  969. let h2handler = NIOHTTP2Handler(
  970. mode: .client,
  971. eventLoop: channel.eventLoop,
  972. streamDelegate: idleHandler
  973. ) { channel in
  974. channel.eventLoop.makeSucceededVoidFuture()
  975. }
  976. XCTAssertNoThrow(try channel.pipeline.syncOperations.addHandler(h2handler))
  977. idleHandler.setMultiplexer(try! h2handler.syncMultiplexer())
  978. self.channels.append(.init(channel: channel, streamDelegate: idleHandler, isActive: false))
  979. XCTAssertNoThrow(try channel.pipeline.syncOperations.addHandler(idleHandler))
  980. return eventLoop.makeSucceededFuture(channel)
  981. }
  982. }
  983. internal struct HookedStreamLender: StreamLender {
  984. internal var onReturnStreams: (Int) -> Void
  985. internal var onUpdateMaxAvailableStreams: (Int) -> Void
  986. internal func returnStreams(_ count: Int, to pool: ConnectionPool) {
  987. self.onReturnStreams(count)
  988. }
  989. internal func changeStreamCapacity(by delta: Int, for: ConnectionPool) {
  990. self.onUpdateMaxAvailableStreams(delta)
  991. }
  992. }
  993. extension Optional where Wrapped == ConnectionPoolError {
  994. internal var isTooManyWaiters: Bool {
  995. switch self {
  996. case .some(.tooManyWaiters):
  997. return true
  998. case .some(.deadlineExceeded), .some(.shutdown), .none:
  999. return false
  1000. }
  1001. }
  1002. internal var isDeadlineExceeded: Bool {
  1003. switch self {
  1004. case .some(.deadlineExceeded):
  1005. return true
  1006. case .some(.tooManyWaiters), .some(.shutdown), .none:
  1007. return false
  1008. }
  1009. }
  1010. internal var isShutdown: Bool {
  1011. switch self {
  1012. case .some(.shutdown):
  1013. return true
  1014. case .some(.tooManyWaiters), .some(.deadlineExceeded), .none:
  1015. return false
  1016. }
  1017. }
  1018. }
  1019. // Simplified version of the frame encoder found in SwiftNIO HTTP/2
  1020. struct HTTP2FrameEncoder {
  1021. mutating func encode(frame: HTTP2Frame, to buf: inout ByteBuffer) throws -> IOData? {
  1022. // note our starting point
  1023. let start = buf.writerIndex
  1024. // +-----------------------------------------------+
  1025. // | Length (24) |
  1026. // +---------------+---------------+---------------+
  1027. // | Type (8) | Flags (8) |
  1028. // +-+-------------+---------------+-------------------------------+
  1029. // |R| Stream Identifier (31) |
  1030. // +=+=============================================================+
  1031. // | Frame Payload (0...) ...
  1032. // +---------------------------------------------------------------+
  1033. // skip 24-bit length for now, we'll fill that in later
  1034. buf.moveWriterIndex(forwardBy: 3)
  1035. // 8-bit type
  1036. buf.writeInteger(frame.code())
  1037. // skip the 8 bit flags for now, we'll fill it in later as well.
  1038. let flagsIndex = buf.writerIndex
  1039. var flags = FrameFlags()
  1040. buf.moveWriterIndex(forwardBy: 1)
  1041. // 32-bit stream identifier -- ensuring the top bit is empty
  1042. buf.writeInteger(Int32(frame.streamID))
  1043. // frame payload follows, which depends on the frame type itself
  1044. let extraFrameData: IOData?
  1045. let payloadSize: Int
  1046. switch frame.payload {
  1047. case let .settings(.settings(settings)):
  1048. for setting in settings {
  1049. buf.writeInteger(setting.parameter.networkRepresentation())
  1050. buf.writeInteger(UInt32(setting.value))
  1051. }
  1052. payloadSize = settings.count * 6
  1053. extraFrameData = nil
  1054. case .settings(.ack):
  1055. payloadSize = 0
  1056. extraFrameData = nil
  1057. flags.insert(.ack)
  1058. case let .goAway(lastStreamID, errorCode, opaqueData):
  1059. let streamVal = UInt32(Int(lastStreamID)) & ~0x8000_0000
  1060. buf.writeInteger(streamVal)
  1061. buf.writeInteger(UInt32(errorCode.networkCode))
  1062. if let data = opaqueData {
  1063. payloadSize = data.readableBytes + 8
  1064. extraFrameData = .byteBuffer(data)
  1065. } else {
  1066. payloadSize = 8
  1067. extraFrameData = nil
  1068. }
  1069. case .data, .headers, .priority,
  1070. .rstStream, .pushPromise, .ping,
  1071. .windowUpdate, .alternativeService, .origin:
  1072. preconditionFailure("Frame type not supported: \(frame.payload)")
  1073. }
  1074. // Write the frame data. This is the payload size and the flags byte.
  1075. buf.writePayloadSize(payloadSize, at: start)
  1076. buf.setInteger(flags.rawValue, at: flagsIndex)
  1077. // all bytes to write are in the provided buffer now
  1078. return extraFrameData
  1079. }
  1080. struct FrameFlags: OptionSet {
  1081. internal private(set) var rawValue: UInt8
  1082. internal init(rawValue: UInt8) {
  1083. self.rawValue = rawValue
  1084. }
  1085. /// ACK flag. Valid on SETTINGS and PING frames.
  1086. internal static let ack = FrameFlags(rawValue: 0x01)
  1087. }
  1088. }
  1089. extension HTTP2SettingsParameter {
  1090. internal func networkRepresentation() -> UInt16 {
  1091. switch self {
  1092. case HTTP2SettingsParameter.headerTableSize:
  1093. return UInt16(1)
  1094. case HTTP2SettingsParameter.enablePush:
  1095. return UInt16(2)
  1096. case HTTP2SettingsParameter.maxConcurrentStreams:
  1097. return UInt16(3)
  1098. case HTTP2SettingsParameter.initialWindowSize:
  1099. return UInt16(4)
  1100. case HTTP2SettingsParameter.maxFrameSize:
  1101. return UInt16(5)
  1102. case HTTP2SettingsParameter.maxHeaderListSize:
  1103. return UInt16(6)
  1104. case HTTP2SettingsParameter.enableConnectProtocol:
  1105. return UInt16(8)
  1106. default:
  1107. preconditionFailure("Unknown settings parameter.")
  1108. }
  1109. }
  1110. }
  1111. extension ByteBuffer {
  1112. fileprivate mutating func writePayloadSize(_ size: Int, at location: Int) {
  1113. // Yes, this performs better than running a UInt8 through the generic write(integer:) three times.
  1114. var bytes: (UInt8, UInt8, UInt8)
  1115. bytes.0 = UInt8((size & 0xFF0000) >> 16)
  1116. bytes.1 = UInt8((size & 0x00FF00) >> 8)
  1117. bytes.2 = UInt8(size & 0x0000FF)
  1118. withUnsafeBytes(of: bytes) { ptr in
  1119. _ = self.setBytes(ptr, at: location)
  1120. }
  1121. }
  1122. }
  1123. extension HTTP2Frame {
  1124. internal func encode() throws -> ByteBuffer {
  1125. let allocator = ByteBufferAllocator()
  1126. var buffer = allocator.buffer(capacity: 1024)
  1127. var frameEncoder = HTTP2FrameEncoder()
  1128. let extraData = try frameEncoder.encode(frame: self, to: &buffer)
  1129. if let extraData = extraData {
  1130. switch extraData {
  1131. case let .byteBuffer(extraBuffer):
  1132. buffer.writeImmutableBuffer(extraBuffer)
  1133. default:
  1134. preconditionFailure()
  1135. }
  1136. }
  1137. return buffer
  1138. }
  1139. /// The one-byte identifier used to indicate the type of a frame on the wire.
  1140. internal func code() -> UInt8 {
  1141. switch self.payload {
  1142. case .data: return 0x0
  1143. case .headers: return 0x1
  1144. case .priority: return 0x2
  1145. case .rstStream: return 0x3
  1146. case .settings: return 0x4
  1147. case .pushPromise: return 0x5
  1148. case .ping: return 0x6
  1149. case .goAway: return 0x7
  1150. case .windowUpdate: return 0x8
  1151. case .alternativeService: return 0xA
  1152. case .origin: return 0xC
  1153. }
  1154. }
  1155. }