Golang Conn.SetWriteDeadline方法代碼示例

func writer(conn net.Conn, outbox chan []byte) {
	defer func() {
		conn.Close()
		clearOutbox(outbox) //incase reader is blocked
	}()

	b := bytes.NewBuffer(nil)

	for {
		select {
		case msg, ok := <-outbox:
			if !ok {
				return
			}
			b.Write(msg)
			for n := len(outbox); n > 0; n-- {
				b.Write(<-outbox)
			}

			conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
			_, err := conn.Write(b.Bytes())
			if err != nil {
				return
			}
			b.Reset()
		}
	}
}

func serveSPDY(conn net.Conn, srv *http.Server) {
	defer func() {
		if v := recover(); v != nil {
			const size = 4096
			buf := make([]byte, size)
			buf = buf[:runtime.Stack(buf, false)]
			log.Printf("panic serving %v: %v\n%s", conn.RemoteAddr(), v, buf)
		}
	}()

	tlsConn, ok := conn.(*tls.Conn)
	if !ok { // Only allow TLS connections.
		return
	}

	if d := srv.ReadTimeout; d != 0 {
		conn.SetReadDeadline(time.Now().Add(d))
	}
	if d := srv.WriteTimeout; d != 0 {
		conn.SetWriteDeadline(time.Now().Add(d))
	}
	if err := tlsConn.Handshake(); err != nil {
		return
	}

	tlsState := new(tls.ConnectionState)
	*tlsState = tlsConn.ConnectionState()
	proto := tlsState.NegotiatedProtocol
	if fn := srv.TLSNextProto[proto]; fn != nil {
		fn(srv, tlsConn, nil)
	}
	return
}

func copyBytes(dst, src net.Conn, wg *sync.WaitGroup, writeDeadline, readDeadline time.Duration) {
	defer wg.Done()
	n, err := io.Copy(dst, src)
	if err != nil {
		log.Errorf("proxy i/o error: %v", err)
	}

	if cr, ok := src.(*net.TCPConn); ok {
		cr.CloseRead()
	} else {
		// For TLS connections.
		wto := time.Now().Add(writeDeadline)
		src.SetWriteDeadline(wto)
	}

	if cw, ok := dst.(*net.TCPConn); ok {
		cw.CloseWrite()
	} else {
		// For TLS connections.
		rto := time.Now().Add(readDeadline)
		dst.SetReadDeadline(rto)
	}

	log.Debugf("proxy copied %d bytes %s -> %s", n, src.RemoteAddr(), dst.RemoteAddr())
}

func (s *Server) updateWriteDeadline(c net.Conn, ctx *RequestCtx, lastDeadlineTime time.Time) time.Time {
	writeTimeout := s.WriteTimeout
	if s.MaxKeepaliveDuration > 0 {
		connTimeout := s.MaxKeepaliveDuration - time.Since(ctx.connTime)
		if connTimeout <= 0 {
			// MaxKeepAliveDuration exceeded, but let's try sending response anyway
			// in 100ms with 'Connection: close' header.
			ctx.SetConnectionClose()
			connTimeout = 100 * time.Millisecond
		}
		if connTimeout < writeTimeout {
			writeTimeout = connTimeout
		}
	}

	// Optimization: update write deadline only if more than 25%
	// of the last write deadline exceeded.
	// See https://github.com/golang/go/issues/15133 for details.
	currentTime := time.Now()
	if currentTime.Sub(lastDeadlineTime) > (writeTimeout >> 2) {
		if err := c.SetWriteDeadline(currentTime.Add(writeTimeout)); err != nil {
			panic(fmt.Sprintf("BUG: error in SetWriteDeadline(%s): %s", writeTimeout, err))
		}
		lastDeadlineTime = currentTime
	}
	return lastDeadlineTime
}

func createPeerWriter(conn net.Conn) (chan<- []byte, <-chan error) {
	msgChan := make(chan []byte)
	errChan := make(chan error)

	go func() {
		defer log.V(3).Infof("WINSTON: Peer writer goroutine exited\n")
		defer close(errChan)
		// msgChan should be closed by the caller

		for msg := range msgChan {
			// Set a deadline for sending the next message and refresh it before each message
			conn.SetWriteDeadline(time.Now().Add(30 * time.Second))

			err := netWriteUint32(conn, uint32(len(msg)))
			if err != nil {
				errChan <- fmt.Errorf("Could not send byte of new message: '%s'", err)
				break
			}
			_, err = conn.Write(msg)
			if err != nil {
				errChan <- fmt.Errorf("Could not send a message: '%s'", err)
				break
			}
		}
	}()

	return msgChan, errChan
}

// poll sends a packet and wait for a response. Both operations can timeout, they're retried up to retries times.
func poll(conn net.Conn, toSend []byte, respondBuffer []byte, retries int, timeout time.Duration) (int, error) {
	var err error
	for i := 0; i < retries+1; i++ {
		deadline := time.Now().Add(timeout)

		if err = conn.SetWriteDeadline(deadline); err != nil {
			log.Printf("Couldn't set write deadline. Retrying. Retry %d/%d\n", i, retries)
			continue
		}
		if _, err = conn.Write(toSend); err != nil {
			log.Printf("Couldn't write. Retrying. Retry %d/%d\n", i, retries)
			continue
		}

		deadline = time.Now().Add(timeout)
		if err = conn.SetReadDeadline(deadline); err != nil {
			log.Printf("Couldn't set read deadline. Retrying. Retry %d/%d\n", i, retries)
			continue
		}

		numRead := 0
		if numRead, err = conn.Read(respondBuffer); err != nil {
			log.Printf("Couldn't read. Retrying. Retry %d/%d\n", i, retries)
			continue
		}

		return numRead, nil
	}
	return 0, err
}

func copyContent(from net.Conn, to net.Conn, complete chan bool, done chan bool, otherDone chan bool) {
	var err error = nil
	var bytes []byte = make([]byte, 256)
	var read int = 0
	for {
		select {
		// If we received a done message from the other goroutine, we exit.
		case <-otherDone:
			complete <- true
			return
		default:

			// Read data from the source connection.
			from.SetReadDeadline(time.Now().Add(time.Second * 5))
			read, err = from.Read(bytes)
			// If any errors occured, write to complete as we are done (one of the
			// connections closed.)
			if err != nil {
				complete <- true
				done <- true
				return
			}
			// Write data to the destination.
			to.SetWriteDeadline(time.Now().Add(time.Second * 5))
			_, err = to.Write(bytes[:read])
			// Same error checking.
			if err != nil {
				complete <- true
				done <- true
				return
			}
		}
	}
}

// testRacyWrite tests that it is safe to mutate the input Write buffer
// immediately after cancelation has occurred.
func testRacyWrite(t *testing.T, c1, c2 net.Conn) {
	go chunkedCopy(ioutil.Discard, c2)

	var wg sync.WaitGroup
	defer wg.Wait()

	c1.SetWriteDeadline(time.Now().Add(time.Millisecond))
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			b1 := make([]byte, 1024)
			b2 := make([]byte, 1024)
			for j := 0; j < 100; j++ {
				_, err := c1.Write(b1)
				copy(b1, b2) // Mutate b1 to trigger potential race
				if err != nil {
					checkForTimeoutError(t, err)
					c1.SetWriteDeadline(time.Now().Add(time.Millisecond))
				}
			}
		}()
	}
}

func (service *StreamService) serve(conn net.Conn) {
	var data []byte
	var err error
	for {
		if service.readTimeout != nil {
			err = conn.SetReadDeadline(time.Now().Add(service.readTimeout.(time.Duration)))
		}
		if err == nil {
			data, err = receiveDataOverStream(conn)
		}
		if err == nil {
			data = service.Handle(data, &StreamContext{BaseContext: NewBaseContext(), Conn: conn})
			if service.writeTimeout != nil {
				err = conn.SetWriteDeadline(time.Now().Add(service.writeTimeout.(time.Duration)))
			}
			if err == nil {
				err = sendDataOverStream(conn, data)
			}
		}
		if err != nil {
			conn.Close()
			break
		}
	}
}

// copy Content two-way
func pass(from net.Conn, to net.Conn, complete chan bool, oneSide chan bool, otherSide chan bool) {
	var err error
	var read int
	bytes := make([]byte, 256)

	for {
		select {

		case <-otherSide:
			complete <- true
			return

		default:

			from.SetReadDeadline(time.Now().Add(time.Duration(pConfig.Timeout) * time.Second))
			read, err = from.Read(bytes)
			if err != nil {
				complete <- true
				oneSide <- true
				return
			}

			to.SetWriteDeadline(time.Now().Add(time.Duration(pConfig.Timeout) * time.Second))
			_, err = to.Write(bytes[:read])
			if err != nil {
				complete <- true
				oneSide <- true
				return
			}
		}
	}
}

func serveSPDY(conn net.Conn, srv *http.Server) {
	defer common.Recover()

	tlsConn, ok := conn.(*tls.Conn)
	if !ok { // Only allow TLS connections.
		return
	}

	if d := srv.ReadTimeout; d != 0 {
		conn.SetReadDeadline(time.Now().Add(d))
	}
	if d := srv.WriteTimeout; d != 0 {
		conn.SetWriteDeadline(time.Now().Add(d))
	}
	if err := tlsConn.Handshake(); err != nil {
		return
	}

	tlsState := new(tls.ConnectionState)
	*tlsState = tlsConn.ConnectionState()
	proto := tlsState.NegotiatedProtocol
	if fn := srv.TLSNextProto[proto]; fn != nil {
		fn(srv, tlsConn, nil)
	}
	return
}

func (s *session) proxy(c1, c2 net.Conn) error {
	if debug {
		log.Println("Proxy", c1.RemoteAddr(), "->", c2.RemoteAddr())
	}

	atomic.AddInt64(&numProxies, 1)
	defer atomic.AddInt64(&numProxies, -1)

	buf := make([]byte, 65536)
	for {
		c1.SetReadDeadline(time.Now().Add(networkTimeout))
		n, err := c1.Read(buf)
		if err != nil {
			return err
		}

		atomic.AddInt64(&bytesProxied, int64(n))

		if debug {
			log.Printf("%d bytes from %s to %s", n, c1.RemoteAddr(), c2.RemoteAddr())
		}

		if s.rateLimit != nil {
			s.rateLimit(int64(n))
		}

		c2.SetWriteDeadline(time.Now().Add(networkTimeout))
		_, err = c2.Write(buf[:n])
		if err != nil {
			return err
		}
	}
}

// copy Content two-way
func pass(from net.Conn, to net.Conn, complete chan bool, one_side chan bool, other_side chan bool) {
	var err error = nil
	var bytes []byte = make([]byte, 256)
	var read int = 0

	for {
		select {

		case <-other_side:
			complete <- true
			return

		default:

			from.SetReadDeadline(time.Now().Add(time.Duration(*expire) * time.Second))
			read, err = from.Read(bytes)
			if err != nil {
				complete <- true
				one_side <- true
				return
			}

			to.SetWriteDeadline(time.Now().Add(time.Duration(*expire) * time.Second))
			_, err = to.Write(bytes[:read])
			if err != nil {
				complete <- true
				one_side <- true
				return
			}
		}
	}
}

func WriteMessage(conn net.Conn, timeout time.Duration, msg *Message) error {
	var handle = &codec.MsgpackHandle{}
	var buf []byte
	var encoder = codec.NewEncoderBytes(&buf, handle)
	err := encoder.Encode(msg)
	if err != nil {
		return err
	}

	length := len(buf)
	if length > DefaultMessageLengthCap {
		return errors.New("message is too long")
	}

	var lengthBuf [2]byte
	binary.BigEndian.PutUint16(lengthBuf[:], uint16(length))

	conn.SetWriteDeadline(time.Now().Add(timeout))
	_, err = conn.Write(lengthBuf[:])
	if err != nil {
		return err
	}

	_, err = io.Copy(conn, bytes.NewBuffer(buf))
	return err
}

func write(conn net.Conn, content []byte, timeout float64) error {
	if timeout > 0 {
		conn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
	}
	_, err := conn.Write(content)
	return err
}