Sfoglia il codice sorgente

Decrease an amount of goroutines

tags/0.15^2
9seconds 7 anni fa
parent
commit
76cfbb009a
2 ha cambiato i file con 26 aggiunte e 53 eliminazioni
  1. 6
    16
      proxy/proxy.go
  2. 20
    37
      wrappers/conn.go

+ 6
- 16
proxy/proxy.go Vedi File

93
 		clientPacket := clientConn.(wrappers.PacketReadWriteCloser)
93
 		clientPacket := clientConn.(wrappers.PacketReadWriteCloser)
94
 		serverPacket := serverConn.(wrappers.PacketReadWriteCloser)
94
 		serverPacket := serverConn.(wrappers.PacketReadWriteCloser)
95
 		go p.middlePipe(clientPacket, serverPacket, wait, &opts.ReadHacks)
95
 		go p.middlePipe(clientPacket, serverPacket, wait, &opts.ReadHacks)
96
-		go p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
96
+		p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
97
 	} else {
97
 	} else {
98
 		clientStream := clientConn.(wrappers.StreamReadWriteCloser)
98
 		clientStream := clientConn.(wrappers.StreamReadWriteCloser)
99
 		serverStream := serverConn.(wrappers.StreamReadWriteCloser)
99
 		serverStream := serverConn.(wrappers.StreamReadWriteCloser)
100
 		go p.directPipe(clientStream, serverStream, wait, p.conf.ReadBufferSize)
100
 		go p.directPipe(clientStream, serverStream, wait, p.conf.ReadBufferSize)
101
-		go p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
101
+		p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
102
 	}
102
 	}
103
 
103
 
104
 	wait.Wait()
104
 	wait.Wait()
121
 	return packetConn, nil
121
 	return packetConn, nil
122
 }
122
 }
123
 
123
 
124
-func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser,
125
-	wait *sync.WaitGroup, hacks *mtproto.Hacks) {
126
-	defer func() {
127
-		src.Close() // nolint: errcheck, gosec
128
-		dst.Close() // nolint: errcheck, gosec
129
-		wait.Done()
130
-	}()
124
+func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.Writer, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
125
+	defer wait.Done()
131
 
126
 
132
 	for {
127
 	for {
133
 		hacks.SimpleAck = false
128
 		hacks.SimpleAck = false
145
 	}
140
 	}
146
 }
141
 }
147
 
142
 
148
-func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.WriteCloser,
149
-	wait *sync.WaitGroup, bufferSize int) {
150
-	defer func() {
151
-		src.Close() // nolint: errcheck, gosec
152
-		dst.Close() // nolint: errcheck, gosec
153
-		wait.Done()
154
-	}()
143
+func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.Writer, wait *sync.WaitGroup, bufferSize int) {
144
+	defer wait.Done()
155
 
145
 
156
 	buffer := make([]byte, bufferSize)
146
 	buffer := make([]byte, bufferSize)
157
 	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {
147
 	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {

+ 20
- 37
wrappers/conn.go Vedi File

38
 	connTimeoutWrite = 2 * time.Minute
38
 	connTimeoutWrite = 2 * time.Minute
39
 )
39
 )
40
 
40
 
41
-type ioResult struct {
42
-	n   int
43
-	err error
44
-}
45
-
46
-type ioFunc func([]byte) (int, error)
47
-
48
 // Conn is a basic wrapper for net.Conn providing the most low-level
41
 // Conn is a basic wrapper for net.Conn providing the most low-level
49
 // logic and management as possible.
42
 // logic and management as possible.
50
 type Conn struct {
43
 type Conn struct {
61
 func (c *Conn) Write(p []byte) (int, error) {
54
 func (c *Conn) Write(p []byte) (int, error) {
62
 	select {
55
 	select {
63
 	case <-c.ctx.Done():
56
 	case <-c.ctx.Done():
57
+		c.Close() // nolint: gosec
64
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot write because context was closed")
58
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot write because context was closed")
65
 	default:
59
 	default:
66
-		n, err := c.doIO(c.conn.Write, p, connTimeoutWrite)
60
+		if err := c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)); err != nil {
61
+			c.Close() // nolint: gosec
62
+			return 0, errors.Annotate(err, "Cannot set write deadline to the socket")
63
+		}
67
 
64
 
65
+		n, err := c.conn.Write(p)
68
 		c.logger.Debugw("Write to stream", "bytes", n, "error", err)
66
 		c.logger.Debugw("Write to stream", "bytes", n, "error", err)
69
 		stats.EgressTraffic(n)
67
 		stats.EgressTraffic(n)
68
+		if err != nil {
69
+			c.Close() // nolint: gosec
70
+		}
70
 
71
 
71
 		return n, err
72
 		return n, err
72
 	}
73
 	}
75
 func (c *Conn) Read(p []byte) (int, error) {
76
 func (c *Conn) Read(p []byte) (int, error) {
76
 	select {
77
 	select {
77
 	case <-c.ctx.Done():
78
 	case <-c.ctx.Done():
79
+		c.Close() // nolint: gosec
78
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot read because context was closed")
80
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot read because context was closed")
79
 	default:
81
 	default:
80
-		n, err := c.doIO(c.conn.Read, p, connTimeoutRead)
82
+		if err := c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)); err != nil {
83
+			c.Close() // nolint: gosec
84
+			return 0, errors.Annotate(err, "Cannot set read deadline to the socket")
85
+		}
81
 
86
 
87
+		n, err := c.conn.Read(p)
82
 		c.logger.Debugw("Read from stream", "bytes", n, "error", err)
88
 		c.logger.Debugw("Read from stream", "bytes", n, "error", err)
83
 		stats.IngressTraffic(n)
89
 		stats.IngressTraffic(n)
84
-
85
-		return n, err
86
-	}
87
-}
88
-
89
-func (c *Conn) doIO(callback ioFunc, p []byte, timeout time.Duration) (int, error) {
90
-	resChan := make(chan ioResult, 1)
91
-	timer := time.NewTimer(timeout)
92
-
93
-	go func() {
94
-		n, err := callback(p)
95
-		resChan <- ioResult{n: n, err: err}
96
-	}()
97
-
98
-	select {
99
-	case res := <-resChan:
100
-		timer.Stop()
101
-		if res.err != nil {
90
+		if err != nil {
102
 			c.Close() // nolint: gosec
91
 			c.Close() // nolint: gosec
103
 		}
92
 		}
104
-		return res.n, res.err
105
-	case <-c.ctx.Done():
106
-		timer.Stop()
107
-		c.Close() // nolint: gosec
108
-		return 0, errors.Annotate(c.ctx.Err(), "Cannot do IO because context is closed")
109
-	case <-timer.C:
110
-		c.Close() // nolint: gosec
111
-		return 0, errors.Annotate(c.ctx.Err(), "Timeout on IO operation")
93
+
94
+		return n, err
112
 	}
95
 	}
113
 }
96
 }
114
 
97
 
115
 // Close closes underlying net.Conn instance.
98
 // Close closes underlying net.Conn instance.
116
 func (c *Conn) Close() error {
99
 func (c *Conn) Close() error {
117
-	defer c.logger.Debugw("Close connection")
118
-
100
+	c.logger.Debugw("Close connection")
119
 	c.cancel()
101
 	c.cancel()
102
+
120
 	return c.conn.Close()
103
 	return c.conn.Close()
121
 }
104
 }
122
 
105
 

Loading…
Annulla
Salva