CommandLineToolkit

1
/*
2
 * Copyright (c) Avito Tech LLC
3
 */
4

5
import Foundation
6
import Socket
7
import SocketModels
8

9
// `LazySocketConnection` will connect socket on demand and reconnect it in case of link failure.
10
public final class LazySocketConnection {
11
    public typealias SocketFactory = () throws -> Socket
12

13
    private let socketAddress: SocketAddress
14
    private let socketFactory: SocketFactory
15

16
    private var lazySocket: Socket?
17
    private var socket: Socket {
18
        get throws {
19
            if let socket = lazySocket {
20
                return socket
21
            }
22

23
            let newSocket = try makeConnectedSocket()
24
            lazySocket = newSocket
25
            return newSocket
26
        }
27
    }
28

29
    public init(
30
        socketAddress: SocketAddress,
31
        socketFactory: @escaping SocketFactory
32
    ) {
33
        self.socketAddress = socketAddress
34
        self.socketFactory = socketFactory
35
    }
36

37
    public func send(data: Data, retriesLimit: Int = 1) throws {
38
        try send(data: data, tryNumber: 0, retriesLimit: retriesLimit)
39
    }
40

41
    public func close() {
42
        lazySocket?.close()
43
        lazySocket = nil
44
    }
45

46
    private func send(data: Data, tryNumber: Int, retriesLimit: Int) throws {
47
        do {
48
            try socket.write(from: data)
49
        } catch {
50
            lazySocket = nil
51

52
            if tryNumber < retriesLimit {
53
                try send(data: data, tryNumber: tryNumber + 1, retriesLimit: retriesLimit)
54
            } else {
55
                throw error
56
            }
57
        }
58
    }
59

60
    private func makeConnectedSocket() throws -> Socket {
61
        let socket = try socketFactory()
62
        try socket.connect(
63
            to: socketAddress.host,
64
            port: Int32(socketAddress.port.value)
65
        )
66
        return socket
67
    }
68
}
69