Merge pull request #44 from 9seconds/goroutines-number

Decrease the number of used goroutines

Dockerfile

@@ -10,8 +10,7 @@ RUN set -x \
     curl \
     git \
     make \
-    upx \
-  && update-ca-certificates
+    upx
 
 COPY . /go/src/github.com/9seconds/mtg/
 
@@ -26,7 +25,7 @@ RUN set -x \
 
 FROM scratch
 
-ENTRYPOINT ["/usr/local/bin/mtg"]
+ENTRYPOINT ["/mtg"]
 ENV MTG_IP=0.0.0.0 \
     MTG_PORT=3128 \
     MTG_STATS_IP=0.0.0.0 \
@@ -34,4 +33,4 @@ ENV MTG_IP=0.0.0.0 \
 EXPOSE 3128 3129
 
 COPY --from=0 /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
-COPY --from=0 /go/src/github.com/9seconds/mtg/mtg /usr/local/bin/mtg
+COPY --from=0 /go/src/github.com/9seconds/mtg/mtg /mtg

go.mod

@@ -7,18 +7,17 @@ require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dustin/go-humanize v1.0.0
 	github.com/gofrs/uuid v3.1.0+incompatible
-	github.com/juju/errors v0.0.0-20180806074554-22422dad46e1
+	github.com/juju/errors v0.0.0-20181012004132-a4583d0a56ea
 	github.com/juju/loggo v0.0.0-20180524022052-584905176618 // indirect
 	github.com/juju/testing v0.0.0-20180920084828-472a3e8b2073 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
-	github.com/kr/pty v1.1.3 // indirect
 	github.com/pkg/errors v0.8.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stretchr/testify v1.2.2
 	go.uber.org/atomic v1.3.2 // indirect
 	go.uber.org/multierr v1.1.0 // indirect
 	go.uber.org/zap v1.9.1
-	golang.org/x/net v0.0.0-20181005035420-146acd28ed58 // indirect
+	golang.org/x/net v0.0.0-20181011144130-49bb7cea24b1 // indirect
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6
 	gopkg.in/alexcesaro/statsd.v2 v2.0.0
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect

go.sum

@@ -6,14 +6,12 @@ github.com/beevik/ntp v0.2.0 h1:sGsd+kAXzT0bfVfzJfce04g+dSRfrs+tbQW8lweuYgw=
 github.com/beevik/ntp v0.2.0/go.mod h1:hIHWr+l3+/clUnF44zdK+CWW7fO8dR5cIylAQ76NRpg=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/dustin/go-humanize v0.0.0-20180713052910-9f541cc9db5d h1:lDrio3iIdNb0Gw9CgH7cQF+iuB5mOOjdJ9ERNJCBgb4=
-github.com/dustin/go-humanize v0.0.0-20180713052910-9f541cc9db5d/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/gofrs/uuid v3.1.0+incompatible h1:q2rtkjaKT4YEr6E1kamy0Ha4RtepWlQBedyHx0uzKwA=
 github.com/gofrs/uuid v3.1.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
-github.com/juju/errors v0.0.0-20180806074554-22422dad46e1 h1:wnhMXidtb70kDZCeLt/EfsVtkXS5c8zLnE9y/6DIRAU=
-github.com/juju/errors v0.0.0-20180806074554-22422dad46e1/go.mod h1:W54LbzXuIE0boCoNJfwqpmkKJ1O4TCTZMetAt6jGk7Q=
+github.com/juju/errors v0.0.0-20181012004132-a4583d0a56ea h1:g2k+8WR7cHch4g0tBDhfiEvAp7fXxTNBiD1oC1Oxj3E=
+github.com/juju/errors v0.0.0-20181012004132-a4583d0a56ea/go.mod h1:W54LbzXuIE0boCoNJfwqpmkKJ1O4TCTZMetAt6jGk7Q=
 github.com/juju/loggo v0.0.0-20180524022052-584905176618 h1:MK144iBQF9hTSwBW/9eJm034bVoG30IshVm688T2hi8=
 github.com/juju/loggo v0.0.0-20180524022052-584905176618/go.mod h1:vgyd7OREkbtVEN/8IXZe5Ooef3LQePvuBm9UWj6ZL8U=
 github.com/juju/testing v0.0.0-20180920084828-472a3e8b2073 h1:WQM1NildKThwdP7qWrNAFGzp4ijNLw8RlgENkaI4MJs=
@@ -21,7 +19,6 @@ github.com/juju/testing v0.0.0-20180920084828-472a3e8b2073/go.mod h1:63prj8cnj0t
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
-github.com/kr/pty v1.1.3/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/pkg/errors v0.8.0 h1:WdK/asTD0HN+q6hsWO3/vpuAkAr+tw6aNJNDFFf0+qw=
@@ -36,10 +33,8 @@ go.uber.org/multierr v1.1.0 h1:HoEmRHQPVSqub6w2z2d2EOVs2fjyFRGyofhKuyDq0QI=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/zap v1.9.1 h1:XCJQEf3W6eZaVwhRBof6ImoYGJSITeKWsyeh3HFu/5o=
 go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
-golang.org/x/net v0.0.0-20180921000356-2f5d2388922f h1:QM2QVxvDoW9PFSPp/zy9FgxJLfaWTZlS61KEPtBwacM=
-golang.org/x/net v0.0.0-20180921000356-2f5d2388922f/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
-golang.org/x/net v0.0.0-20181005035420-146acd28ed58 h1:otZG8yDCO4LVps5+9bxOeNiCvgmOyt96J3roHTYs7oE=
-golang.org/x/net v0.0.0-20181005035420-146acd28ed58/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+golang.org/x/net v0.0.0-20181011144130-49bb7cea24b1 h1:Y/KGZSOdz/2r0WJ9Mkmz6NJBusp0kiNx1Cn82lzJQ6w=
+golang.org/x/net v0.0.0-20181011144130-49bb7cea24b1/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6 h1:jMFz6MfLP0/4fUyZle81rXUoxOBFi19VUFKVDOQfozc=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/alexcesaro/statsd.v2 v2.0.0 h1:FXkZSCZIH17vLCO5sO2UucTHsH9pc+17F6pl3JVCwMc=

proxy/proxy.go

@@ -93,12 +93,12 @@ func (p *Proxy) accept(conn net.Conn) {
 		clientPacket := clientConn.(wrappers.PacketReadWriteCloser)
 		serverPacket := serverConn.(wrappers.PacketReadWriteCloser)
 		go p.middlePipe(clientPacket, serverPacket, wait, &opts.ReadHacks)
-		go p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
+		p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
 	} else {
 		clientStream := clientConn.(wrappers.StreamReadWriteCloser)
 		serverStream := serverConn.(wrappers.StreamReadWriteCloser)
 		go p.directPipe(clientStream, serverStream, wait, p.conf.ReadBufferSize)
-		go p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
+		p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
 	}
 
 	wait.Wait()
@@ -121,13 +121,8 @@ func (p *Proxy) getTelegramConn(ctx context.Context, cancel context.CancelFunc,
 	return packetConn, nil
 }
 
-func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser,
-	wait *sync.WaitGroup, hacks *mtproto.Hacks) {
-	defer func() {
-		src.Close() // nolint: errcheck, gosec
-		dst.Close() // nolint: errcheck, gosec
-		wait.Done()
-	}()
+func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.Writer, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
+	defer wait.Done()
 
 	for {
 		hacks.SimpleAck = false
@@ -145,13 +140,8 @@ func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser,
 	}
 }
 
-func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.WriteCloser,
-	wait *sync.WaitGroup, bufferSize int) {
-	defer func() {
-		src.Close() // nolint: errcheck, gosec
-		dst.Close() // nolint: errcheck, gosec
-		wait.Done()
-	}()
+func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.Writer, wait *sync.WaitGroup, bufferSize int) {
+	defer wait.Done()
 
 	buffer := make([]byte, bufferSize)
 	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {

wrappers/conn.go

@@ -38,13 +38,6 @@ const (
 	connTimeoutWrite = 2 * time.Minute
 )
 
-type ioResult struct {
-	n   int
-	err error
-}
-
-type ioFunc func([]byte) (int, error)
-
 // Conn is a basic wrapper for net.Conn providing the most low-level
 // logic and management as possible.
 type Conn struct {
@@ -61,12 +54,20 @@ type Conn struct {
 func (c *Conn) Write(p []byte) (int, error) {
 	select {
 	case <-c.ctx.Done():
+		c.Close() // nolint: gosec
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot write because context was closed")
 	default:
-		n, err := c.doIO(c.conn.Write, p, connTimeoutWrite)
+		if err := c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)); err != nil {
+			c.Close() // nolint: gosec
+			return 0, errors.Annotate(err, "Cannot set write deadline to the socket")
+		}
 
+		n, err := c.conn.Write(p)
 		c.logger.Debugw("Write to stream", "bytes", n, "error", err)
 		stats.EgressTraffic(n)
+		if err != nil {
+			c.Close() // nolint: gosec
+		}
 
 		return n, err
 	}
@@ -75,48 +76,30 @@ func (c *Conn) Write(p []byte) (int, error) {
 func (c *Conn) Read(p []byte) (int, error) {
 	select {
 	case <-c.ctx.Done():
+		c.Close() // nolint: gosec
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot read because context was closed")
 	default:
-		n, err := c.doIO(c.conn.Read, p, connTimeoutRead)
+		if err := c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)); err != nil {
+			c.Close() // nolint: gosec
+			return 0, errors.Annotate(err, "Cannot set read deadline to the socket")
+		}
 
+		n, err := c.conn.Read(p)
 		c.logger.Debugw("Read from stream", "bytes", n, "error", err)
 		stats.IngressTraffic(n)
-
-		return n, err
-	}
-}
-
-func (c *Conn) doIO(callback ioFunc, p []byte, timeout time.Duration) (int, error) {
-	resChan := make(chan ioResult, 1)
-	timer := time.NewTimer(timeout)
-
-	go func() {
-		n, err := callback(p)
-		resChan <- ioResult{n: n, err: err}
-	}()
-
-	select {
-	case res := <-resChan:
-		timer.Stop()
-		if res.err != nil {
+		if err != nil {
 			c.Close() // nolint: gosec
 		}
-		return res.n, res.err
-	case <-c.ctx.Done():
-		timer.Stop()
-		c.Close() // nolint: gosec
-		return 0, errors.Annotate(c.ctx.Err(), "Cannot do IO because context is closed")
-	case <-timer.C:
-		c.Close() // nolint: gosec
-		return 0, errors.Annotate(c.ctx.Err(), "Timeout on IO operation")
+
+		return n, err
 	}
 }
 
 // Close closes underlying net.Conn instance.
 func (c *Conn) Close() error {
-	defer c.logger.Debugw("Close connection")
-
+	c.logger.Debugw("Close connection")
 	c.cancel()
+
 	return c.conn.Close()
 }