<commit_before>package runnerpool

import (
	"context"

	"github.com/fnproject/fn/api/common"
	"github.com/fnproject/fn/api/models"

	"go.opencensus.io/stats"
)

type placerTracker struct {
	cfg        *PlacerConfig
	requestCtx context.Context
	placerCtx  context.Context
	cancel     context.CancelFunc
	tracker    *attemptTracker
	isPlaced   bool
}

func NewPlacerTracker(requestCtx context.Context, cfg *PlacerConfig, call RunnerCall) *placerTracker {

	timeout := cfg.PlacerTimeout
	if call.Model().Type == models.TypeDetached {
		timeout = cfg.DetachedPlacerTimeout
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	return &placerTracker{
		cfg:        cfg,
		requestCtx: requestCtx,
		placerCtx:  ctx,
		cancel:     cancel,
		tracker:    newAttemptTracker(requestCtx),
	}
}

// IsDone is a non-blocking check to see if the underlying deadlines are exceeded.
func (tr *placerTracker) IsDone() bool {
	return tr.requestCtx.Err() != nil || tr.placerCtx.Err() != nil
}

// HandleFindRunnersFailure is a convenience function to record error from runnerpool.Runners()
func (tr *placerTracker) HandleFindRunnersFailure(err error) {
	logger := common.Logger(tr.requestCtx).WithError(err)
	w, ok := err.(models.APIErrorWrapper)
	if ok {
		logger = logger.WithField("root_error", w.RootError())
	}
	logger.Error("Failed to find runners for call")
	stats.Record(tr.requestCtx, errorPoolCountMeasure.M(0))
}

// TryRunner is a convenience function to TryExec a call on a runner and
// analyze the results.
func (tr *placerTracker) TryRunner(r Runner, call RunnerCall) (bool, error) {
	tr.tracker.recordAttempt()

	// WARNING: Do not use placerCtx here to let requestCtx take its time
	// during container execution.
	ctx, cancel := context.WithCancel(tr.requestCtx)
	isPlaced, err := r.TryExec(ctx, call)
	cancel()

	if !isPlaced {

		// Too Busy is super common case, we track it separately
		if err == models.ErrCallTimeoutServerBusy {
			stats.Record(tr.requestCtx, retryTooBusyCountMeasure.M(0))
		} else if tr.requestCtx.Err() != err {
			// only record retry due to an error if client did not abort/cancel/timeout
			stats.Record(tr.requestCtx, retryErrorCountMeasure.M(0))
		}

	} else {
		if err == nil {
			stats.Record(tr.requestCtx, placedOKCountMeasure.M(0))
		} else if tr.requestCtx.Err() == err {
			stats.Record(tr.requestCtx, placedAbortCountMeasure.M(0))
		} else {
			stats.Record(tr.requestCtx, placedErrorCountMeasure.M(0))
		}

		// Call is now committed. In other words, it was 'run'. We are done.
		tr.isPlaced = true
	}

	return isPlaced, err
}

// HandleDone is cleanup function to cancel pending contexts and to
// record stats for the placement session.
func (tr *placerTracker) HandleDone() {

	// Cancel Exit Path / Client cancelled/timedout
	if tr.requestCtx.Err() == context.Canceled {
		stats.Record(tr.requestCtx, cancelCountMeasure.M(0))
	} else if tr.requestCtx.Err() == context.DeadlineExceeded {
		stats.Record(tr.requestCtx, timeoutCountMeasure.M(0))
	}

	// This means our placer timed out. We ignore tr.isPlaced calls
	// since we do not check/track placer ctx timeout if a call was
	// actually ran on a runner. This means, placer timeout can be
	// 10 secs, but a call can execute for 60 secs in a container.
	if !tr.isPlaced && tr.placerCtx.Err() != nil {
		stats.Record(tr.requestCtx, placerTimeoutMeasure.M(0))
	}

	tr.tracker.finalizeAttempts(tr.isPlaced)
	tr.cancel()
}

// RetryAllBackoff blocks until it is time to try the runner list again. Returns
// false if the placer should stop trying.
func (tr *placerTracker) RetryAllBackoff(numOfRunners int, err error) bool {
	// If there are no runners and last call to provision runners failed due
	// to a user error (or misconfiguration) then fail fast.
	if numOfRunners == 0 && err != nil {
		// IsFuncError currently synonymous with tag: 'blame == user'
		// See: runner_fninvoke.handleFnInvokeCall2
		if models.IsFuncError(err) {
			// We also sanity check for a 502 before returning.
			w, ok := err.(models.APIError)
			if ok {
				if 502 == w.Code() {
					return false
				}
			}
		}
	}

	// This means Placer is operating on an empty list. No runners
	// available. Record it.
	if numOfRunners == 0 {
		stats.Record(tr.requestCtx, emptyPoolCountMeasure.M(0))
	}

	t := common.NewTimer(tr.cfg.RetryAllDelay)
	defer t.Stop()

	select {
	case <-tr.requestCtx.Done(): // client side timeout/cancel
		return false
	case <-tr.placerCtx.Done(): // placer wait timeout
		return false
	case <-t.C:
	}

	return true
}
<commit_msg>Decrease log level of failure to find runners message to WARN (#1562)<commit_after>package runnerpool

import (
	"context"

	"github.com/fnproject/fn/api/common"
	"github.com/fnproject/fn/api/models"

	"go.opencensus.io/stats"
)

type placerTracker struct {
	cfg        *PlacerConfig
	requestCtx context.Context
	placerCtx  context.Context
	cancel     context.CancelFunc
	tracker    *attemptTracker
	isPlaced   bool
}

func NewPlacerTracker(requestCtx context.Context, cfg *PlacerConfig, call RunnerCall) *placerTracker {

	timeout := cfg.PlacerTimeout
	if call.Model().Type == models.TypeDetached {
		timeout = cfg.DetachedPlacerTimeout
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	return &placerTracker{
		cfg:        cfg,
		requestCtx: requestCtx,
		placerCtx:  ctx,
		cancel:     cancel,
		tracker:    newAttemptTracker(requestCtx),
	}
}

// IsDone is a non-blocking check to see if the underlying deadlines are exceeded.
func (tr *placerTracker) IsDone() bool {
	return tr.requestCtx.Err() != nil || tr.placerCtx.Err() != nil
}

// HandleFindRunnersFailure is a convenience function to record error from runnerpool.Runners()
func (tr *placerTracker) HandleFindRunnersFailure(err error) {
	logger := common.Logger(tr.requestCtx).WithError(err)
	w, ok := err.(models.APIErrorWrapper)
	if ok {
		logger = logger.WithField("root_error", w.RootError())
	}
	logger.Warn("Failed to find runners for call")
	stats.Record(tr.requestCtx, errorPoolCountMeasure.M(0))
}

// TryRunner is a convenience function to TryExec a call on a runner and
// analyze the results.
func (tr *placerTracker) TryRunner(r Runner, call RunnerCall) (bool, error) {
	tr.tracker.recordAttempt()

	// WARNING: Do not use placerCtx here to let requestCtx take its time
	// during container execution.
	ctx, cancel := context.WithCancel(tr.requestCtx)
	isPlaced, err := r.TryExec(ctx, call)
	cancel()

	if !isPlaced {

		// Too Busy is super common case, we track it separately
		if err == models.ErrCallTimeoutServerBusy {
			stats.Record(tr.requestCtx, retryTooBusyCountMeasure.M(0))
		} else if tr.requestCtx.Err() != err {
			// only record retry due to an error if client did not abort/cancel/timeout
			stats.Record(tr.requestCtx, retryErrorCountMeasure.M(0))
		}

	} else {
		if err == nil {
			stats.Record(tr.requestCtx, placedOKCountMeasure.M(0))
		} else if tr.requestCtx.Err() == err {
			stats.Record(tr.requestCtx, placedAbortCountMeasure.M(0))
		} else {
			stats.Record(tr.requestCtx, placedErrorCountMeasure.M(0))
		}

		// Call is now committed. In other words, it was 'run'. We are done.
		tr.isPlaced = true
	}

	return isPlaced, err
}

// HandleDone is cleanup function to cancel pending contexts and to
// record stats for the placement session.
func (tr *placerTracker) HandleDone() {

	// Cancel Exit Path / Client cancelled/timedout
	if tr.requestCtx.Err() == context.Canceled {
		stats.Record(tr.requestCtx, cancelCountMeasure.M(0))
	} else if tr.requestCtx.Err() == context.DeadlineExceeded {
		stats.Record(tr.requestCtx, timeoutCountMeasure.M(0))
	}

	// This means our placer timed out. We ignore tr.isPlaced calls
	// since we do not check/track placer ctx timeout if a call was
	// actually ran on a runner. This means, placer timeout can be
	// 10 secs, but a call can execute for 60 secs in a container.
	if !tr.isPlaced && tr.placerCtx.Err() != nil {
		stats.Record(tr.requestCtx, placerTimeoutMeasure.M(0))
	}

	tr.tracker.finalizeAttempts(tr.isPlaced)
	tr.cancel()
}

// RetryAllBackoff blocks until it is time to try the runner list again. Returns
// false if the placer should stop trying.
func (tr *placerTracker) RetryAllBackoff(numOfRunners int, err error) bool {
	// If there are no runners and last call to provision runners failed due
	// to a user error (or misconfiguration) then fail fast.
	if numOfRunners == 0 && err != nil {
		// IsFuncError currently synonymous with tag: 'blame == user'
		// See: runner_fninvoke.handleFnInvokeCall2
		if models.IsFuncError(err) {
			// We also sanity check for a 502 before returning.
			w, ok := err.(models.APIError)
			if ok {
				if 502 == w.Code() {
					return false
				}
			}
		}
	}

	// This means Placer is operating on an empty list. No runners
	// available. Record it.
	if numOfRunners == 0 {
		stats.Record(tr.requestCtx, emptyPoolCountMeasure.M(0))
	}

	t := common.NewTimer(tr.cfg.RetryAllDelay)
	defer t.Stop()

	select {
	case <-tr.requestCtx.Done(): // client side timeout/cancel
		return false
	case <-tr.placerCtx.Done(): // placer wait timeout
		return false
	case <-t.C:
	}

	return true
}
<|endoftext|>


<commit_before>/*
 * Copyright 2016, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

// This file is the implementation of a gRPC server using HTTP/2 which
// uses the standard Go http2 Server implementation (via the
// http.Handler interface), rather than speaking low-level HTTP/2
// frames itself. It is the implementation of *grpc.Server.ServeHTTP.

package transport

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"
)

// NewServerHandlerTransport returns a ServerTransport handling gRPC
// from inside an http.Handler. It requires that the http Server
// supports HTTP/2.
func NewServerHandlerTransport(w http.ResponseWriter, r *http.Request) (ServerTransport, error) {
	if r.ProtoMajor != 2 {
		return nil, errors.New("gRPC requires HTTP/2")
	}
	if r.Method != "POST" {
		return nil, errors.New("invalid gRPC request method")
	}
	if !validContentType(r.Header.Get("Content-Type")) {
		return nil, errors.New("invalid gRPC request content-type")
	}
	if _, ok := w.(http.Flusher); !ok {
		return nil, errors.New("gRPC requires a ResponseWriter supporting http.Flusher")
	}
	if _, ok := w.(http.CloseNotifier); !ok {
		return nil, errors.New("gRPC requires a ResponseWriter supporting http.CloseNotifier")
	}

	st := &serverHandlerTransport{
		rw:       w,
		req:      r,
		closedCh: make(chan struct{}),
		writes:   make(chan func()),
	}

	if v := r.Header.Get("grpc-timeout"); v != "" {
		to, err := decodeTimeout(v)
		if err != nil {
			return nil, streamErrorf(codes.Internal, "malformed time-out: %v", err)
		}
		st.timeoutSet = true
		st.timeout = to
	}

	var metakv []string
	if r.Host != "" {
		metakv = append(metakv, ":authority", r.Host)
	}
	for k, vv := range r.Header {
		k = strings.ToLower(k)
		if isReservedHeader(k) && !isWhitelistedPseudoHeader(k) {
			continue
		}
		for _, v := range vv {
			if k == "user-agent" {
				// user-agent is special. Copying logic of http_util.go.
				if i := strings.LastIndex(v, " "); i == -1 {
					// There is no application user agent string being set
					continue
				} else {
					v = v[:i]
				}
			}
			v, err := decodeMetadataHeader(k, v)
			if err != nil {
				return nil, streamErrorf(codes.InvalidArgument, "malformed binary metadata: %v", err)
			}
			metakv = append(metakv, k, v)
		}
	}
	st.headerMD = metadata.Pairs(metakv...)

	return st, nil
}

// serverHandlerTransport is an implementation of ServerTransport
// which replies to exactly one gRPC request (exactly one HTTP request),
// using the net/http.Handler interface. This http.Handler is guaranteed
// at this point to be speaking over HTTP/2, so it's able to speak valid
// gRPC.
type serverHandlerTransport struct {
	rw               http.ResponseWriter
	req              *http.Request
	timeoutSet       bool
	timeout          time.Duration
	didCommonHeaders bool

	headerMD metadata.MD

	closeOnce sync.Once
	closedCh  chan struct{} // closed on Close

	// writes is a channel of code to run serialized in the
	// ServeHTTP (HandleStreams) goroutine. The channel is closed
	// when WriteStatus is called.
	writes chan func()
}

func (ht *serverHandlerTransport) Close() error {
	ht.closeOnce.Do(ht.closeCloseChanOnce)
	return nil
}

func (ht *serverHandlerTransport) closeCloseChanOnce() { close(ht.closedCh) }

func (ht *serverHandlerTransport) RemoteAddr() net.Addr { return strAddr(ht.req.RemoteAddr) }

// strAddr is a net.Addr backed by either a TCP "ip:port" string, or
// the empty string if unknown.
type strAddr string

func (a strAddr) Network() string {
	if a != "" {
		// Per the documentation on net/http.Request.RemoteAddr, if this is
		// set, it's set to the IP:port of the peer (hence, TCP):
		// https://golang.org/pkg/net/http/#Request
		//
		// If we want to support Unix sockets later, we can
		// add our own grpc-specific convention within the
		// grpc codebase to set RemoteAddr to a different
		// format, or probably better: we can attach it to the
		// context and use that from serverHandlerTransport.RemoteAddr.
		return "tcp"
	}
	return ""
}

func (a strAddr) String() string { return string(a) }

// do runs fn in the ServeHTTP goroutine.
func (ht *serverHandlerTransport) do(fn func()) error {
	// Avoid a panic writing to closed channel. Imperfect but maybe good enough.
	select {
	case <-ht.closedCh:
		return ErrConnClosing
	default:
		select {
		case ht.writes <- fn:
			return nil
		case <-ht.closedCh:
			return ErrConnClosing
		}

	}
}

func (ht *serverHandlerTransport) WriteStatus(s *Stream, st *status.Status) error {
	err := ht.do(func() {
		ht.writeCommonHeaders(s)

		// And flush, in case no header or body has been sent yet.
		// This forces a separation of headers and trailers if this is the
		// first call (for example, in end2end tests's TestNoService).
		ht.rw.(http.Flusher).Flush()

		h := ht.rw.Header()
		h.Set("Grpc-Status", fmt.Sprintf("%d", st.Code()))
		if m := st.Message(); m != "" {
			h.Set("Grpc-Message", encodeGrpcMessage(m))
		}

		// TODO: Support Grpc-Status-Details-Bin

		if md := s.Trailer(); len(md) > 0 {
			for k, vv := range md {
				// Clients don't tolerate reading restricted headers after some non restricted ones were sent.
				if isReservedHeader(k) {
					continue
				}
				for _, v := range vv {
					// http2 ResponseWriter mechanism to send undeclared Trailers after
					// the headers have possibly been written.
					h.Add(http2.TrailerPrefix+k, encodeMetadataHeader(k, v))
				}
			}
		}
	})
	close(ht.writes)
	return err
}

// writeCommonHeaders sets common headers on the first write
// call (Write, WriteHeader, or WriteStatus).
func (ht *serverHandlerTransport) writeCommonHeaders(s *Stream) {
	if ht.didCommonHeaders {
		return
	}
	ht.didCommonHeaders = true

	h := ht.rw.Header()
	h["Date"] = nil // suppress Date to make tests happy; TODO: restore
	h.Set("Content-Type", "application/grpc")

	// Predeclare trailers we'll set later in WriteStatus (after the body).
	// This is a SHOULD in the HTTP RFC, and the way you add (known)
	// Trailers per the net/http.ResponseWriter contract.
	// See https://golang.org/pkg/net/http/#ResponseWriter
	// and https://golang.org/pkg/net/http/#example_ResponseWriter_trailers
	h.Add("Trailer", "Grpc-Status")
	h.Add("Trailer", "Grpc-Message")
	// TODO: Support Grpc-Status-Details-Bin

	if s.sendCompress != "" {
		h.Set("Grpc-Encoding", s.sendCompress)
	}
}

func (ht *serverHandlerTransport) Write(s *Stream, data []byte, opts *Options) error {
	return ht.do(func() {
		ht.writeCommonHeaders(s)
		ht.rw.Write(data)
		if !opts.Delay {
			ht.rw.(http.Flusher).Flush()
		}
	})
}

func (ht *serverHandlerTransport) WriteHeader(s *Stream, md metadata.MD) error {
	return ht.do(func() {
		ht.writeCommonHeaders(s)
		h := ht.rw.Header()
		for k, vv := range md {
			// Clients don't tolerate reading restricted headers after some non restricted ones were sent.
			if isReservedHeader(k) {
				continue
			}
			for _, v := range vv {
				v = encodeMetadataHeader(k, v)
				h.Add(k, v)
			}
		}
		ht.rw.WriteHeader(200)
		ht.rw.(http.Flusher).Flush()
	})
}

func (ht *serverHandlerTransport) HandleStreams(startStream func(*Stream), traceCtx func(context.Context, string) context.Context) {
	// With this transport type there will be exactly 1 stream: this HTTP request.

	var ctx context.Context
	var cancel context.CancelFunc
	if ht.timeoutSet {
		ctx, cancel = context.WithTimeout(context.Background(), ht.timeout)
	} else {
		ctx, cancel = context.WithCancel(context.Background())
	}

	// requestOver is closed when either the request's context is done
	// or the status has been written via WriteStatus.
	requestOver := make(chan struct{})

	// clientGone receives a single value if peer is gone, either
	// because the underlying connection is dead or because the
	// peer sends an http2 RST_STREAM.
	clientGone := ht.rw.(http.CloseNotifier).CloseNotify()
	go func() {
		select {
		case <-requestOver:
			return
		case <-ht.closedCh:
		case <-clientGone:
		}
		cancel()
	}()

	req := ht.req

	s := &Stream{
		id:           0, // irrelevant
		cancel:       cancel,
		buf:          newRecvBuffer(),
		st:           ht,
		method:       req.URL.Path,
		recvCompress: req.Header.Get("grpc-encoding"),
	}
	pr := &peer.Peer{
		Addr: ht.RemoteAddr(),
	}
	if req.TLS != nil {
		pr.AuthInfo = credentials.TLSInfo{State: *req.TLS}
	}
	ctx = metadata.NewIncomingContext(ctx, ht.headerMD)
	ctx = peer.NewContext(ctx, pr)
	s.ctx = newContextWithStream(ctx, s)
	s.trReader = &recvBufferReader{ctx: s.ctx, recv: s.buf}

	// readerDone is closed when the Body.Read-ing goroutine exits.
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)

		// TODO: minimize garbage, optimize recvBuffer code/ownership
		const readSize = 8196
		for buf := make([]byte, readSize); ; {
			n, err := req.Body.Read(buf)
			if n > 0 {
				s.buf.put(&recvMsg{data: buf[:n:n]})
				buf = buf[n:]
			}
			if err != nil {
				s.buf.put(&recvMsg{err: mapRecvMsgError(err)})
				return
			}
			if len(buf) == 0 {
				buf = make([]byte, readSize)
			}
		}
	}()

	// startStream is provided by the *grpc.Server's serveStreams.
	// It starts a goroutine serving s and exits immediately.
	// The goroutine that is started is the one that then calls
	// into ht, calling WriteHeader, Write, WriteStatus, Close, etc.
	startStream(s)

	ht.runStream()
	close(requestOver)

	// Wait for reading goroutine to finish.
	req.Body.Close()
	<-readerDone
}

func (ht *serverHandlerTransport) runStream() {
	for {
		select {
		case fn, ok := <-ht.writes:
			if !ok {
				return
			}
			fn()
		case <-ht.closedCh:
			return
		}
	}
}

func (ht *serverHandlerTransport) Drain() {
	panic("Drain() is not implemented")
}

// mapRecvMsgError returns the non-nil err into the appropriate
// error value as expected by callers of *grpc.parser.recvMsg.
// In particular, in can only be:
//   * io.EOF
//   * io.ErrUnexpectedEOF
//   * of type transport.ConnectionError
//   * of type transport.StreamError
func mapRecvMsgError(err error) error {
	if err == io.EOF || err == io.ErrUnexpectedEOF {
		return err
	}
	if se, ok := err.(http2.StreamError); ok {
		if code, ok := http2ErrConvTab[se.Code]; ok {
			return StreamError{
				Code: code,
				Desc: se.Error(),
			}
		}
	}
	return connectionErrorf(true, err, err.Error())
}
<commit_msg>Initialize stream properly in handler_server. (#1260)<commit_after>/*
 * Copyright 2016, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

// This file is the implementation of a gRPC server using HTTP/2 which
// uses the standard Go http2 Server implementation (via the
// http.Handler interface), rather than speaking low-level HTTP/2
// frames itself. It is the implementation of *grpc.Server.ServeHTTP.

package transport

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"
)

// NewServerHandlerTransport returns a ServerTransport handling gRPC
// from inside an http.Handler. It requires that the http Server
// supports HTTP/2.
func NewServerHandlerTransport(w http.ResponseWriter, r *http.Request) (ServerTransport, error) {
	if r.ProtoMajor != 2 {
		return nil, errors.New("gRPC requires HTTP/2")
	}
	if r.Method != "POST" {
		return nil, errors.New("invalid gRPC request method")
	}
	if !validContentType(r.Header.Get("Content-Type")) {
		return nil, errors.New("invalid gRPC request content-type")
	}
	if _, ok := w.(http.Flusher); !ok {
		return nil, errors.New("gRPC requires a ResponseWriter supporting http.Flusher")
	}
	if _, ok := w.(http.CloseNotifier); !ok {
		return nil, errors.New("gRPC requires a ResponseWriter supporting http.CloseNotifier")
	}

	st := &serverHandlerTransport{
		rw:       w,
		req:      r,
		closedCh: make(chan struct{}),
		writes:   make(chan func()),
	}

	if v := r.Header.Get("grpc-timeout"); v != "" {
		to, err := decodeTimeout(v)
		if err != nil {
			return nil, streamErrorf(codes.Internal, "malformed time-out: %v", err)
		}
		st.timeoutSet = true
		st.timeout = to
	}

	var metakv []string
	if r.Host != "" {
		metakv = append(metakv, ":authority", r.Host)
	}
	for k, vv := range r.Header {
		k = strings.ToLower(k)
		if isReservedHeader(k) && !isWhitelistedPseudoHeader(k) {
			continue
		}
		for _, v := range vv {
			if k == "user-agent" {
				// user-agent is special. Copying logic of http_util.go.
				if i := strings.LastIndex(v, " "); i == -1 {
					// There is no application user agent string being set
					continue
				} else {
					v = v[:i]
				}
			}
			v, err := decodeMetadataHeader(k, v)
			if err != nil {
				return nil, streamErrorf(codes.InvalidArgument, "malformed binary metadata: %v", err)
			}
			metakv = append(metakv, k, v)
		}
	}
	st.headerMD = metadata.Pairs(metakv...)

	return st, nil
}

// serverHandlerTransport is an implementation of ServerTransport
// which replies to exactly one gRPC request (exactly one HTTP request),
// using the net/http.Handler interface. This http.Handler is guaranteed
// at this point to be speaking over HTTP/2, so it's able to speak valid
// gRPC.
type serverHandlerTransport struct {
	rw               http.ResponseWriter
	req              *http.Request
	timeoutSet       bool
	timeout          time.Duration
	didCommonHeaders bool

	headerMD metadata.MD

	closeOnce sync.Once
	closedCh  chan struct{} // closed on Close

	// writes is a channel of code to run serialized in the
	// ServeHTTP (HandleStreams) goroutine. The channel is closed
	// when WriteStatus is called.
	writes chan func()
}

func (ht *serverHandlerTransport) Close() error {
	ht.closeOnce.Do(ht.closeCloseChanOnce)
	return nil
}

func (ht *serverHandlerTransport) closeCloseChanOnce() { close(ht.closedCh) }

func (ht *serverHandlerTransport) RemoteAddr() net.Addr { return strAddr(ht.req.RemoteAddr) }

// strAddr is a net.Addr backed by either a TCP "ip:port" string, or
// the empty string if unknown.
type strAddr string

func (a strAddr) Network() string {
	if a != "" {
		// Per the documentation on net/http.Request.RemoteAddr, if this is
		// set, it's set to the IP:port of the peer (hence, TCP):
		// https://golang.org/pkg/net/http/#Request
		//
		// If we want to support Unix sockets later, we can
		// add our own grpc-specific convention within the
		// grpc codebase to set RemoteAddr to a different
		// format, or probably better: we can attach it to the
		// context and use that from serverHandlerTransport.RemoteAddr.
		return "tcp"
	}
	return ""
}

func (a strAddr) String() string { return string(a) }

// do runs fn in the ServeHTTP goroutine.
func (ht *serverHandlerTransport) do(fn func()) error {
	// Avoid a panic writing to closed channel. Imperfect but maybe good enough.
	select {
	case <-ht.closedCh:
		return ErrConnClosing
	default:
		select {
		case ht.writes <- fn:
			return nil
		case <-ht.closedCh:
			return ErrConnClosing
		}

	}
}

func (ht *serverHandlerTransport) WriteStatus(s *Stream, st *status.Status) error {
	err := ht.do(func() {
		ht.writeCommonHeaders(s)

		// And flush, in case no header or body has been sent yet.
		// This forces a separation of headers and trailers if this is the
		// first call (for example, in end2end tests's TestNoService).
		ht.rw.(http.Flusher).Flush()

		h := ht.rw.Header()
		h.Set("Grpc-Status", fmt.Sprintf("%d", st.Code()))
		if m := st.Message(); m != "" {
			h.Set("Grpc-Message", encodeGrpcMessage(m))
		}

		// TODO: Support Grpc-Status-Details-Bin

		if md := s.Trailer(); len(md) > 0 {
			for k, vv := range md {
				// Clients don't tolerate reading restricted headers after some non restricted ones were sent.
				if isReservedHeader(k) {
					continue
				}
				for _, v := range vv {
					// http2 ResponseWriter mechanism to send undeclared Trailers after
					// the headers have possibly been written.
					h.Add(http2.TrailerPrefix+k, encodeMetadataHeader(k, v))
				}
			}
		}
	})
	close(ht.writes)
	return err
}

// writeCommonHeaders sets common headers on the first write
// call (Write, WriteHeader, or WriteStatus).
func (ht *serverHandlerTransport) writeCommonHeaders(s *Stream) {
	if ht.didCommonHeaders {
		return
	}
	ht.didCommonHeaders = true

	h := ht.rw.Header()
	h["Date"] = nil // suppress Date to make tests happy; TODO: restore
	h.Set("Content-Type", "application/grpc")

	// Predeclare trailers we'll set later in WriteStatus (after the body).
	// This is a SHOULD in the HTTP RFC, and the way you add (known)
	// Trailers per the net/http.ResponseWriter contract.
	// See https://golang.org/pkg/net/http/#ResponseWriter
	// and https://golang.org/pkg/net/http/#example_ResponseWriter_trailers
	h.Add("Trailer", "Grpc-Status")
	h.Add("Trailer", "Grpc-Message")
	// TODO: Support Grpc-Status-Details-Bin

	if s.sendCompress != "" {
		h.Set("Grpc-Encoding", s.sendCompress)
	}
}

func (ht *serverHandlerTransport) Write(s *Stream, data []byte, opts *Options) error {
	return ht.do(func() {
		ht.writeCommonHeaders(s)
		ht.rw.Write(data)
		if !opts.Delay {
			ht.rw.(http.Flusher).Flush()
		}
	})
}

func (ht *serverHandlerTransport) WriteHeader(s *Stream, md metadata.MD) error {
	return ht.do(func() {
		ht.writeCommonHeaders(s)
		h := ht.rw.Header()
		for k, vv := range md {
			// Clients don't tolerate reading restricted headers after some non restricted ones were sent.
			if isReservedHeader(k) {
				continue
			}
			for _, v := range vv {
				v = encodeMetadataHeader(k, v)
				h.Add(k, v)
			}
		}
		ht.rw.WriteHeader(200)
		ht.rw.(http.Flusher).Flush()
	})
}

func (ht *serverHandlerTransport) HandleStreams(startStream func(*Stream), traceCtx func(context.Context, string) context.Context) {
	// With this transport type there will be exactly 1 stream: this HTTP request.

	var ctx context.Context
	var cancel context.CancelFunc
	if ht.timeoutSet {
		ctx, cancel = context.WithTimeout(context.Background(), ht.timeout)
	} else {
		ctx, cancel = context.WithCancel(context.Background())
	}

	// requestOver is closed when either the request's context is done
	// or the status has been written via WriteStatus.
	requestOver := make(chan struct{})

	// clientGone receives a single value if peer is gone, either
	// because the underlying connection is dead or because the
	// peer sends an http2 RST_STREAM.
	clientGone := ht.rw.(http.CloseNotifier).CloseNotify()
	go func() {
		select {
		case <-requestOver:
			return
		case <-ht.closedCh:
		case <-clientGone:
		}
		cancel()
	}()

	req := ht.req

	s := &Stream{
		id:           0, // irrelevant
		requestRead:  func(int) {},
		cancel:       cancel,
		buf:          newRecvBuffer(),
		st:           ht,
		method:       req.URL.Path,
		recvCompress: req.Header.Get("grpc-encoding"),
	}
	pr := &peer.Peer{
		Addr: ht.RemoteAddr(),
	}
	if req.TLS != nil {
		pr.AuthInfo = credentials.TLSInfo{State: *req.TLS}
	}
	ctx = metadata.NewIncomingContext(ctx, ht.headerMD)
	ctx = peer.NewContext(ctx, pr)
	s.ctx = newContextWithStream(ctx, s)
	s.trReader = &transportReader{
		reader:        &recvBufferReader{ctx: s.ctx, recv: s.buf},
		windowHandler: func(int) {},
	}

	// readerDone is closed when the Body.Read-ing goroutine exits.
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)

		// TODO: minimize garbage, optimize recvBuffer code/ownership
		const readSize = 8196
		for buf := make([]byte, readSize); ; {
			n, err := req.Body.Read(buf)
			if n > 0 {
				s.buf.put(&recvMsg{data: buf[:n:n]})
				buf = buf[n:]
			}
			if err != nil {
				s.buf.put(&recvMsg{err: mapRecvMsgError(err)})
				return
			}
			if len(buf) == 0 {
				buf = make([]byte, readSize)
			}
		}
	}()

	// startStream is provided by the *grpc.Server's serveStreams.
	// It starts a goroutine serving s and exits immediately.
	// The goroutine that is started is the one that then calls
	// into ht, calling WriteHeader, Write, WriteStatus, Close, etc.
	startStream(s)

	ht.runStream()
	close(requestOver)

	// Wait for reading goroutine to finish.
	req.Body.Close()
	<-readerDone
}

func (ht *serverHandlerTransport) runStream() {
	for {
		select {
		case fn, ok := <-ht.writes:
			if !ok {
				return
			}
			fn()
		case <-ht.closedCh:
			return
		}
	}
}

func (ht *serverHandlerTransport) Drain() {
	panic("Drain() is not implemented")
}

// mapRecvMsgError returns the non-nil err into the appropriate
// error value as expected by callers of *grpc.parser.recvMsg.
// In particular, in can only be:
//   * io.EOF
//   * io.ErrUnexpectedEOF
//   * of type transport.ConnectionError
//   * of type transport.StreamError
func mapRecvMsgError(err error) error {
	if err == io.EOF || err == io.ErrUnexpectedEOF {
		return err
	}
	if se, ok := err.(http2.StreamError); ok {
		if code, ok := http2ErrConvTab[se.Code]; ok {
			return StreamError{
				Code: code,
				Desc: se.Error(),
			}
		}
	}
	return connectionErrorf(true, err, err.Error())
}
<|endoftext|>


<commit_before>package transport

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	maddr "github.com/micro/util/go/lib/addr"
	mnet "github.com/micro/util/go/lib/net"
	mls "github.com/micro/util/go/lib/tls"
	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"
)

type buffer struct {
	io.ReadWriter
}

type httpTransport struct {
	opts Options
}

type httpTransportClient struct {
	ht       *httpTransport
	addr     string
	conn     net.Conn
	dialOpts DialOptions
	once     sync.Once

	sync.RWMutex
	r    chan *http.Request
	bl   []*http.Request
	buff *bufio.Reader

	// local/remote ip
	local  string
	remote string
}

type httpTransportSocket struct {
	ht *httpTransport
	w  http.ResponseWriter
	r  *http.Request
	rw *bufio.ReadWriter

	conn net.Conn
	// for the first request
	ch chan *http.Request

	// local/remote ip
	local  string
	remote string
}

type httpTransportListener struct {
	ht       *httpTransport
	listener net.Listener
}

func (b *buffer) Close() error {
	return nil
}

func (h *httpTransportClient) Local() string {
	return h.local
}

func (h *httpTransportClient) Remote() string {
	return h.remote
}

func (h *httpTransportClient) Send(m *Message) error {
	header := make(http.Header)

	for k, v := range m.Header {
		header.Set(k, v)
	}

	reqB := bytes.NewBuffer(m.Body)
	defer reqB.Reset()
	buf := &buffer{
		reqB,
	}

	req := &http.Request{
		Method: "POST",
		URL: &url.URL{
			Scheme: "http",
			Host:   h.addr,
		},
		Header:        header,
		Body:          buf,
		ContentLength: int64(reqB.Len()),
		Host:          h.addr,
	}

	h.Lock()
	h.bl = append(h.bl, req)
	select {
	case h.r <- h.bl[0]:
		h.bl = h.bl[1:]
	default:
	}
	h.Unlock()

	// set timeout if its greater than 0
	if h.ht.opts.Timeout > time.Duration(0) {
		h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
	}

	return req.Write(h.conn)
}

func (h *httpTransportClient) Recv(m *Message) error {
	if m == nil {
		return errors.New("message passed in is nil")
	}

	var r *http.Request
	if !h.dialOpts.Stream {
		rc, ok := <-h.r
		if !ok {
			return io.EOF
		}
		r = rc
	}

	h.RLock()
	if h.buff == nil {
		return io.EOF
	}
	h.RUnlock()

	// set timeout if its greater than 0
	if h.ht.opts.Timeout > time.Duration(0) {
		h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
	}

	rsp, err := http.ReadResponse(h.buff, r)
	if err != nil {
		return err
	}
	defer rsp.Body.Close()

	b, err := ioutil.ReadAll(rsp.Body)
	if err != nil {
		return err
	}

	if rsp.StatusCode != 200 {
		return errors.New(rsp.Status + ": " + string(b))
	}

	m.Body = b

	if m.Header == nil {
		m.Header = make(map[string]string)
	}

	for k, v := range rsp.Header {
		if len(v) > 0 {
			m.Header[k] = v[0]
		} else {
			m.Header[k] = ""
		}
	}

	return nil
}

func (h *httpTransportClient) Close() error {
	err := h.conn.Close()
	h.once.Do(func() {
		h.Lock()
		h.buff.Reset(nil)
		h.buff = nil
		h.Unlock()
		close(h.r)
	})
	return err
}

func (h *httpTransportSocket) Local() string {
	return h.local
}

func (h *httpTransportSocket) Remote() string {
	return h.remote
}

func (h *httpTransportSocket) Recv(m *Message) error {
	if m == nil {
		return errors.New("message passed in is nil")
	}

	if m.Header == nil {
		m.Header = make(map[string]string)
	}

	// process http 1
	if h.r.ProtoMajor == 1 {
		// set timeout if its greater than 0
		if h.ht.opts.Timeout > time.Duration(0) {
			h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
		}

		var r *http.Request

		select {
		// get first request
		case r = <-h.ch:
		// read next request
		default:
			rr, err := http.ReadRequest(h.rw.Reader)
			if err != nil {
				return err
			}
			r = rr
		}

		// read body
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return err
		}

		// set body
		r.Body.Close()
		m.Body = b

		// set headers
		for k, v := range r.Header {
			if len(v) > 0 {
				m.Header[k] = v[0]
			} else {
				m.Header[k] = ""
			}
		}

		// return early early
		return nil
	}

	// processing http2 request
	// read streaming body

	// set max buffer size
	buf := make([]byte, 4*1024)

	// read the request body
	n, err := h.r.Body.Read(buf)
	// not an eof error
	if err != nil {
		return err
	}

	// check if we have data
	if n > 0 {
		m.Body = buf[:n]
	}

	// set headers
	for k, v := range h.r.Header {
		if len(v) > 0 {
			m.Header[k] = v[0]
		} else {
			m.Header[k] = ""
		}
	}

	// set path
	m.Header[":path"] = h.r.URL.Path

	return nil
}

func (h *httpTransportSocket) Send(m *Message) error {
	if h.r.ProtoMajor == 1 {
		rsp := &http.Response{
			Header:        h.r.Header,
			Body:          ioutil.NopCloser(bytes.NewReader(m.Body)),
			Status:        "200 OK",
			StatusCode:    200,
			Proto:         "HTTP/1.1",
			ProtoMajor:    1,
			ProtoMinor:    1,
			ContentLength: int64(len(m.Body)),
		}

		for k, v := range m.Header {
			rsp.Header.Set(k, v)
		}

		// set timeout if its greater than 0
		if h.ht.opts.Timeout > time.Duration(0) {
			h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
		}

		return rsp.Write(h.conn)
	}

	// http2 request

	// set headers
	for k, v := range m.Header {
		h.w.Header().Set(k, v)
	}

	// write request
	_, err := h.w.Write(m.Body)
	return err
}

func (h *httpTransportSocket) error(m *Message) error {
	if h.r.ProtoMajor == 1 {
		rsp := &http.Response{
			Header:        make(http.Header),
			Body:          ioutil.NopCloser(bytes.NewReader(m.Body)),
			Status:        "500 Internal Server Error",
			StatusCode:    500,
			Proto:         "HTTP/1.1",
			ProtoMajor:    1,
			ProtoMinor:    1,
			ContentLength: int64(len(m.Body)),
		}

		for k, v := range m.Header {
			rsp.Header.Set(k, v)
		}

		return rsp.Write(h.conn)
	}
	return nil
}

func (h *httpTransportSocket) Close() error {
	if h.r.ProtoMajor == 1 {
		return h.conn.Close()
	}
	return nil
}

func (h *httpTransportListener) Addr() string {
	return h.listener.Addr().String()
}

func (h *httpTransportListener) Close() error {
	return h.listener.Close()
}

func (h *httpTransportListener) Accept(fn func(Socket)) error {
	// create handler mux
	mux := http.NewServeMux()

	// register our transport handler
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var buf *bufio.ReadWriter
		var con net.Conn

		// read a regular request
		if r.ProtoMajor == 1 {
			b, err := ioutil.ReadAll(r.Body)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			r.Body = ioutil.NopCloser(bytes.NewReader(b))
			// hijack the conn
			hj, ok := w.(http.Hijacker)
			if !ok {
				// we're screwed
				http.Error(w, "cannot serve conn", http.StatusInternalServerError)
				return
			}

			conn, bufrw, err := hj.Hijack()
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			defer conn.Close()
			buf = bufrw
			con = conn
		}

		// save the request
		ch := make(chan *http.Request, 1)
		ch <- r

		fn(&httpTransportSocket{
			ht:     h.ht,
			w:      w,
			r:      r,
			rw:     buf,
			ch:     ch,
			conn:   con,
			local:  h.Addr(),
			remote: r.RemoteAddr,
		})
	})

	// get optional handlers
	if h.ht.opts.Context != nil {
		handlers, ok := h.ht.opts.Context.Value("http_handlers").(map[string]http.Handler)
		if ok {
			for pattern, handler := range handlers {
				mux.Handle(pattern, handler)
			}
		}
	}

	// default http2 server
	srv := &http.Server{
		Handler: mux,
	}

	// insecure connection use h2c
	if !(h.ht.opts.Secure || h.ht.opts.TLSConfig != nil) {
		srv.Handler = h2c.NewHandler(mux, &http2.Server{})
	}

	// begin serving
	return srv.Serve(h.listener)
}

func (h *httpTransport) Dial(addr string, opts ...DialOption) (Client, error) {
	dopts := DialOptions{
		Timeout: DefaultDialTimeout,
	}

	for _, opt := range opts {
		opt(&dopts)
	}

	var conn net.Conn
	var err error

	// TODO: support dial option here rather than using internal config
	if h.opts.Secure || h.opts.TLSConfig != nil {
		config := h.opts.TLSConfig
		if config == nil {
			config = &tls.Config{
				InsecureSkipVerify: true,
			}
		}
		conn, err = newConn(func(addr string) (net.Conn, error) {
			return tls.DialWithDialer(&net.Dialer{Timeout: dopts.Timeout}, "tcp", addr, config)
		})(addr)
	} else {
		conn, err = newConn(func(addr string) (net.Conn, error) {
			return net.DialTimeout("tcp", addr, dopts.Timeout)
		})(addr)
	}

	if err != nil {
		return nil, err
	}

	return &httpTransportClient{
		ht:       h,
		addr:     addr,
		conn:     conn,
		buff:     bufio.NewReader(conn),
		dialOpts: dopts,
		r:        make(chan *http.Request, 1),
		local:    conn.LocalAddr().String(),
		remote:   conn.RemoteAddr().String(),
	}, nil
}

func (h *httpTransport) Listen(addr string, opts ...ListenOption) (Listener, error) {
	var options ListenOptions
	for _, o := range opts {
		o(&options)
	}

	var l net.Listener
	var err error

	// TODO: support use of listen options
	if h.opts.Secure || h.opts.TLSConfig != nil {
		config := h.opts.TLSConfig

		fn := func(addr string) (net.Listener, error) {
			if config == nil {
				hosts := []string{addr}

				// check if its a valid host:port
				if host, _, err := net.SplitHostPort(addr); err == nil {
					if len(host) == 0 {
						hosts = maddr.IPs()
					} else {
						hosts = []string{host}
					}
				}

				// generate a certificate
				cert, err := mls.Certificate(hosts...)
				if err != nil {
					return nil, err
				}
				config = &tls.Config{Certificates: []tls.Certificate{cert}}
			}
			return tls.Listen("tcp", addr, config)
		}

		l, err = mnet.Listen(addr, fn)
	} else {
		fn := func(addr string) (net.Listener, error) {
			return net.Listen("tcp", addr)
		}

		l, err = mnet.Listen(addr, fn)
	}

	if err != nil {
		return nil, err
	}

	return &httpTransportListener{
		ht:       h,
		listener: l,
	}, nil
}

func (h *httpTransport) Init(opts ...Option) error {
	for _, o := range opts {
		o(&h.opts)
	}
	return nil
}

func (h *httpTransport) Options() Options {
	return h.opts
}

func (h *httpTransport) String() string {
	return "http"
}

func newHTTPTransport(opts ...Option) *httpTransport {
	var options Options
	for _, o := range opts {
		o(&options)
	}
	return &httpTransport{opts: options}
}
<commit_msg>fix race in http transport<commit_after>package transport

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	maddr "github.com/micro/util/go/lib/addr"
	mnet "github.com/micro/util/go/lib/net"
	mls "github.com/micro/util/go/lib/tls"
	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"
)

type buffer struct {
	io.ReadWriter
}

type httpTransport struct {
	opts Options
}

type httpTransportClient struct {
	ht       *httpTransport
	addr     string
	conn     net.Conn
	dialOpts DialOptions
	once     sync.Once

	sync.RWMutex
	r    chan *http.Request
	bl   []*http.Request
	buff *bufio.Reader

	// local/remote ip
	local  string
	remote string
}

type httpTransportSocket struct {
	ht *httpTransport
	w  http.ResponseWriter
	r  *http.Request
	rw *bufio.ReadWriter

	conn net.Conn
	// for the first request
	ch chan *http.Request

	// local/remote ip
	local  string
	remote string
}

type httpTransportListener struct {
	ht       *httpTransport
	listener net.Listener
}

func (b *buffer) Close() error {
	return nil
}

func (h *httpTransportClient) Local() string {
	return h.local
}

func (h *httpTransportClient) Remote() string {
	return h.remote
}

func (h *httpTransportClient) Send(m *Message) error {
	header := make(http.Header)

	for k, v := range m.Header {
		header.Set(k, v)
	}

	reqB := bytes.NewBuffer(m.Body)
	defer reqB.Reset()
	buf := &buffer{
		reqB,
	}

	req := &http.Request{
		Method: "POST",
		URL: &url.URL{
			Scheme: "http",
			Host:   h.addr,
		},
		Header:        header,
		Body:          buf,
		ContentLength: int64(reqB.Len()),
		Host:          h.addr,
	}

	h.Lock()
	h.bl = append(h.bl, req)
	select {
	case h.r <- h.bl[0]:
		h.bl = h.bl[1:]
	default:
	}
	h.Unlock()

	// set timeout if its greater than 0
	if h.ht.opts.Timeout > time.Duration(0) {
		h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
	}

	return req.Write(h.conn)
}

func (h *httpTransportClient) Recv(m *Message) error {
	if m == nil {
		return errors.New("message passed in is nil")
	}

	var r *http.Request
	if !h.dialOpts.Stream {
		rc, ok := <-h.r
		if !ok {
			return io.EOF
		}
		r = rc
	}

	// set timeout if its greater than 0
	if h.ht.opts.Timeout > time.Duration(0) {
		h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
	}

	rsp, err := http.ReadResponse(h.buff, r)
	if err != nil {
		return err
	}
	defer rsp.Body.Close()

	b, err := ioutil.ReadAll(rsp.Body)
	if err != nil {
		return err
	}

	if rsp.StatusCode != 200 {
		return errors.New(rsp.Status + ": " + string(b))
	}

	m.Body = b

	if m.Header == nil {
		m.Header = make(map[string]string)
	}

	for k, v := range rsp.Header {
		if len(v) > 0 {
			m.Header[k] = v[0]
		} else {
			m.Header[k] = ""
		}
	}

	return nil
}

func (h *httpTransportClient) Close() error {
	err := h.conn.Close()
	h.once.Do(func() {
		h.Lock()
		h.buff.Reset(nil)
		h.Unlock()
		close(h.r)
	})
	return err
}

func (h *httpTransportSocket) Local() string {
	return h.local
}

func (h *httpTransportSocket) Remote() string {
	return h.remote
}

func (h *httpTransportSocket) Recv(m *Message) error {
	if m == nil {
		return errors.New("message passed in is nil")
	}

	if m.Header == nil {
		m.Header = make(map[string]string)
	}

	// process http 1
	if h.r.ProtoMajor == 1 {
		// set timeout if its greater than 0
		if h.ht.opts.Timeout > time.Duration(0) {
			h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
		}

		var r *http.Request

		select {
		// get first request
		case r = <-h.ch:
		// read next request
		default:
			rr, err := http.ReadRequest(h.rw.Reader)
			if err != nil {
				return err
			}
			r = rr
		}

		// read body
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return err
		}

		// set body
		r.Body.Close()
		m.Body = b

		// set headers
		for k, v := range r.Header {
			if len(v) > 0 {
				m.Header[k] = v[0]
			} else {
				m.Header[k] = ""
			}
		}

		// return early early
		return nil
	}

	// processing http2 request
	// read streaming body

	// set max buffer size
	buf := make([]byte, 4*1024)

	// read the request body
	n, err := h.r.Body.Read(buf)
	// not an eof error
	if err != nil {
		return err
	}

	// check if we have data
	if n > 0 {
		m.Body = buf[:n]
	}

	// set headers
	for k, v := range h.r.Header {
		if len(v) > 0 {
			m.Header[k] = v[0]
		} else {
			m.Header[k] = ""
		}
	}

	// set path
	m.Header[":path"] = h.r.URL.Path

	return nil
}

func (h *httpTransportSocket) Send(m *Message) error {
	if h.r.ProtoMajor == 1 {
		rsp := &http.Response{
			Header:        h.r.Header,
			Body:          ioutil.NopCloser(bytes.NewReader(m.Body)),
			Status:        "200 OK",
			StatusCode:    200,
			Proto:         "HTTP/1.1",
			ProtoMajor:    1,
			ProtoMinor:    1,
			ContentLength: int64(len(m.Body)),
		}

		for k, v := range m.Header {
			rsp.Header.Set(k, v)
		}

		// set timeout if its greater than 0
		if h.ht.opts.Timeout > time.Duration(0) {
			h.conn.SetDeadline(time.Now().Add(h.ht.opts.Timeout))
		}

		return rsp.Write(h.conn)
	}

	// http2 request

	// set headers
	for k, v := range m.Header {
		h.w.Header().Set(k, v)
	}

	// write request
	_, err := h.w.Write(m.Body)
	return err
}

func (h *httpTransportSocket) error(m *Message) error {
	if h.r.ProtoMajor == 1 {
		rsp := &http.Response{
			Header:        make(http.Header),
			Body:          ioutil.NopCloser(bytes.NewReader(m.Body)),
			Status:        "500 Internal Server Error",
			StatusCode:    500,
			Proto:         "HTTP/1.1",
			ProtoMajor:    1,
			ProtoMinor:    1,
			ContentLength: int64(len(m.Body)),
		}

		for k, v := range m.Header {
			rsp.Header.Set(k, v)
		}

		return rsp.Write(h.conn)
	}
	return nil
}

func (h *httpTransportSocket) Close() error {
	if h.r.ProtoMajor == 1 {
		return h.conn.Close()
	}
	return nil
}

func (h *httpTransportListener) Addr() string {
	return h.listener.Addr().String()
}

func (h *httpTransportListener) Close() error {
	return h.listener.Close()
}

func (h *httpTransportListener) Accept(fn func(Socket)) error {
	// create handler mux
	mux := http.NewServeMux()

	// register our transport handler
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var buf *bufio.ReadWriter
		var con net.Conn

		// read a regular request
		if r.ProtoMajor == 1 {
			b, err := ioutil.ReadAll(r.Body)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			r.Body = ioutil.NopCloser(bytes.NewReader(b))
			// hijack the conn
			hj, ok := w.(http.Hijacker)
			if !ok {
				// we're screwed
				http.Error(w, "cannot serve conn", http.StatusInternalServerError)
				return
			}

			conn, bufrw, err := hj.Hijack()
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			defer conn.Close()
			buf = bufrw
			con = conn
		}

		// save the request
		ch := make(chan *http.Request, 1)
		ch <- r

		fn(&httpTransportSocket{
			ht:     h.ht,
			w:      w,
			r:      r,
			rw:     buf,
			ch:     ch,
			conn:   con,
			local:  h.Addr(),
			remote: r.RemoteAddr,
		})
	})

	// get optional handlers
	if h.ht.opts.Context != nil {
		handlers, ok := h.ht.opts.Context.Value("http_handlers").(map[string]http.Handler)
		if ok {
			for pattern, handler := range handlers {
				mux.Handle(pattern, handler)
			}
		}
	}

	// default http2 server
	srv := &http.Server{
		Handler: mux,
	}

	// insecure connection use h2c
	if !(h.ht.opts.Secure || h.ht.opts.TLSConfig != nil) {
		srv.Handler = h2c.NewHandler(mux, &http2.Server{})
	}

	// begin serving
	return srv.Serve(h.listener)
}

func (h *httpTransport) Dial(addr string, opts ...DialOption) (Client, error) {
	dopts := DialOptions{
		Timeout: DefaultDialTimeout,
	}

	for _, opt := range opts {
		opt(&dopts)
	}

	var conn net.Conn
	var err error

	// TODO: support dial option here rather than using internal config
	if h.opts.Secure || h.opts.TLSConfig != nil {
		config := h.opts.TLSConfig
		if config == nil {
			config = &tls.Config{
				InsecureSkipVerify: true,
			}
		}
		conn, err = newConn(func(addr string) (net.Conn, error) {
			return tls.DialWithDialer(&net.Dialer{Timeout: dopts.Timeout}, "tcp", addr, config)
		})(addr)
	} else {
		conn, err = newConn(func(addr string) (net.Conn, error) {
			return net.DialTimeout("tcp", addr, dopts.Timeout)
		})(addr)
	}

	if err != nil {
		return nil, err
	}

	return &httpTransportClient{
		ht:       h,
		addr:     addr,
		conn:     conn,
		buff:     bufio.NewReader(conn),
		dialOpts: dopts,
		r:        make(chan *http.Request, 1),
		local:    conn.LocalAddr().String(),
		remote:   conn.RemoteAddr().String(),
	}, nil
}

func (h *httpTransport) Listen(addr string, opts ...ListenOption) (Listener, error) {
	var options ListenOptions
	for _, o := range opts {
		o(&options)
	}

	var l net.Listener
	var err error

	// TODO: support use of listen options
	if h.opts.Secure || h.opts.TLSConfig != nil {
		config := h.opts.TLSConfig

		fn := func(addr string) (net.Listener, error) {
			if config == nil {
				hosts := []string{addr}

				// check if its a valid host:port
				if host, _, err := net.SplitHostPort(addr); err == nil {
					if len(host) == 0 {
						hosts = maddr.IPs()
					} else {
						hosts = []string{host}
					}
				}

				// generate a certificate
				cert, err := mls.Certificate(hosts...)
				if err != nil {
					return nil, err
				}
				config = &tls.Config{Certificates: []tls.Certificate{cert}}
			}
			return tls.Listen("tcp", addr, config)
		}

		l, err = mnet.Listen(addr, fn)
	} else {
		fn := func(addr string) (net.Listener, error) {
			return net.Listen("tcp", addr)
		}

		l, err = mnet.Listen(addr, fn)
	}

	if err != nil {
		return nil, err
	}

	return &httpTransportListener{
		ht:       h,
		listener: l,
	}, nil
}

func (h *httpTransport) Init(opts ...Option) error {
	for _, o := range opts {
		o(&h.opts)
	}
	return nil
}

func (h *httpTransport) Options() Options {
	return h.opts
}

func (h *httpTransport) String() string {
	return "http"
}

func newHTTPTransport(opts ...Option) *httpTransport {
	var options Options
	for _, o := range opts {
		o(&options)
	}
	return &httpTransport{opts: options}
}
<|endoftext|>


<commit_before>package pipe_test

import (
	"errors"
	"io"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/sync/errgroup"

	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	. "v2ray.com/core/transport/pipe"
)

func TestPipeReadWrite(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(1024))

	b := buf.New()
	b.WriteString("abcd")
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b}))

	b2 := buf.New()
	b2.WriteString("efg")
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b2}))

	rb, err := pReader.ReadMultiBuffer()
	common.Must(err)
	if r := cmp.Diff(rb.String(), "abcdefg"); r != "" {
		t.Error(r)
	}
}

func TestPipeInterrupt(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(1024))
	payload := []byte{'a', 'b', 'c', 'd'}
	b := buf.New()
	b.Write(payload)
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b}))
	pWriter.Interrupt()

	rb, err := pReader.ReadMultiBuffer()
	if err != io.ErrClosedPipe {
		t.Fatal("expect io.ErrClosePipe, but got ", err)
	}
	if !rb.IsEmpty() {
		t.Fatal("expect empty buffer, but got ", rb.Len())
	}
}

func TestPipeClose(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(1024))
	payload := []byte{'a', 'b', 'c', 'd'}
	b := buf.New()
	common.Must2(b.Write(payload))
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b}))
	common.Must(pWriter.Close())

	rb, err := pReader.ReadMultiBuffer()
	common.Must(err)
	if rb.String() != string(payload) {
		t.Fatal("expect content ", string(payload), " but actually ", rb.String())
	}

	rb, err = pReader.ReadMultiBuffer()
	if err != io.EOF {
		t.Fatal("expected EOF, but got ", err)
	}
	if !rb.IsEmpty() {
		t.Fatal("expect empty buffer, but got ", rb.String())
	}
}

func TestPipeLimitZero(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(0))
	bb := buf.New()
	common.Must2(bb.Write([]byte{'a', 'b'}))
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{bb}))

	var errg errgroup.Group
	errg.Go(func() error {
		b := buf.New()
		b.Write([]byte{'c', 'd'})
		return pWriter.WriteMultiBuffer(buf.MultiBuffer{b})
	})
	errg.Go(func() error {
		time.Sleep(time.Second)

		var container buf.MultiBufferContainer
		if err := buf.Copy(pReader, &container); err != nil {
			return err
		}

		if r := cmp.Diff(container.String(), "abcd"); r != "" {
			return errors.New(r)
		}
		return nil
	})
	errg.Go(func() error {
		time.Sleep(time.Second * 2)
		return pWriter.Close()
	})
	if err := errg.Wait(); err != nil {
		t.Error(err)
	}
}

func TestPipeWriteMultiThread(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(0))

	var errg errgroup.Group
	for i := 0; i < 10; i++ {
		errg.Go(func() error {
			b := buf.New()
			b.WriteString("abcd")
			return pWriter.WriteMultiBuffer(buf.MultiBuffer{b})
		})
	}
	time.Sleep(time.Millisecond * 100)
	pWriter.Close()
	errg.Wait()

	b, err := pReader.ReadMultiBuffer()
	common.Must(err)
	if r := cmp.Diff(b[0].Bytes(), []byte{'a', 'b', 'c', 'd'}); r != "" {
		t.Error(r)
	}
}

func TestInterfaces(t *testing.T) {
	_ = (buf.Reader)(new(Reader))
	_ = (buf.TimeoutReader)(new(Reader))
}

func BenchmarkPipeReadWrite(b *testing.B) {
	reader, writer := New(WithoutSizeLimit())
	a := buf.New()
	a.Extend(buf.Size)
	c := buf.MultiBuffer{a}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		common.Must(writer.WriteMultiBuffer(c))
		d, err := reader.ReadMultiBuffer()
		common.Must(err)
		c = d
	}
}
<commit_msg>test case for interface<commit_after>package pipe_test

import (
	"errors"
	"io"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/sync/errgroup"

	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	. "v2ray.com/core/transport/pipe"
)

func TestPipeReadWrite(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(1024))

	b := buf.New()
	b.WriteString("abcd")
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b}))

	b2 := buf.New()
	b2.WriteString("efg")
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b2}))

	rb, err := pReader.ReadMultiBuffer()
	common.Must(err)
	if r := cmp.Diff(rb.String(), "abcdefg"); r != "" {
		t.Error(r)
	}
}

func TestPipeInterrupt(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(1024))
	payload := []byte{'a', 'b', 'c', 'd'}
	b := buf.New()
	b.Write(payload)
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b}))
	pWriter.Interrupt()

	rb, err := pReader.ReadMultiBuffer()
	if err != io.ErrClosedPipe {
		t.Fatal("expect io.ErrClosePipe, but got ", err)
	}
	if !rb.IsEmpty() {
		t.Fatal("expect empty buffer, but got ", rb.Len())
	}
}

func TestPipeClose(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(1024))
	payload := []byte{'a', 'b', 'c', 'd'}
	b := buf.New()
	common.Must2(b.Write(payload))
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{b}))
	common.Must(pWriter.Close())

	rb, err := pReader.ReadMultiBuffer()
	common.Must(err)
	if rb.String() != string(payload) {
		t.Fatal("expect content ", string(payload), " but actually ", rb.String())
	}

	rb, err = pReader.ReadMultiBuffer()
	if err != io.EOF {
		t.Fatal("expected EOF, but got ", err)
	}
	if !rb.IsEmpty() {
		t.Fatal("expect empty buffer, but got ", rb.String())
	}
}

func TestPipeLimitZero(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(0))
	bb := buf.New()
	common.Must2(bb.Write([]byte{'a', 'b'}))
	common.Must(pWriter.WriteMultiBuffer(buf.MultiBuffer{bb}))

	var errg errgroup.Group
	errg.Go(func() error {
		b := buf.New()
		b.Write([]byte{'c', 'd'})
		return pWriter.WriteMultiBuffer(buf.MultiBuffer{b})
	})
	errg.Go(func() error {
		time.Sleep(time.Second)

		var container buf.MultiBufferContainer
		if err := buf.Copy(pReader, &container); err != nil {
			return err
		}

		if r := cmp.Diff(container.String(), "abcd"); r != "" {
			return errors.New(r)
		}
		return nil
	})
	errg.Go(func() error {
		time.Sleep(time.Second * 2)
		return pWriter.Close()
	})
	if err := errg.Wait(); err != nil {
		t.Error(err)
	}
}

func TestPipeWriteMultiThread(t *testing.T) {
	pReader, pWriter := New(WithSizeLimit(0))

	var errg errgroup.Group
	for i := 0; i < 10; i++ {
		errg.Go(func() error {
			b := buf.New()
			b.WriteString("abcd")
			return pWriter.WriteMultiBuffer(buf.MultiBuffer{b})
		})
	}
	time.Sleep(time.Millisecond * 100)
	pWriter.Close()
	errg.Wait()

	b, err := pReader.ReadMultiBuffer()
	common.Must(err)
	if r := cmp.Diff(b[0].Bytes(), []byte{'a', 'b', 'c', 'd'}); r != "" {
		t.Error(r)
	}
}

func TestInterfaces(t *testing.T) {
	_ = (buf.Reader)(new(Reader))
	_ = (buf.TimeoutReader)(new(Reader))

	_ = (common.Interruptible)(new(Reader))
	_ = (common.Interruptible)(new(Writer))
	_ = (common.Closable)(new(Writer))
}

func BenchmarkPipeReadWrite(b *testing.B) {
	reader, writer := New(WithoutSizeLimit())
	a := buf.New()
	a.Extend(buf.Size)
	c := buf.MultiBuffer{a}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		common.Must(writer.WriteMultiBuffer(c))
		d, err := reader.ReadMultiBuffer()
		common.Must(err)
		c = d
	}
}
<|endoftext|>


<commit_before>package transport

import (
	"fmt"
	"strconv"
	"testing"
	"time"

	"github.com/stefankopieczek/gossip/base"
)

type endpoint struct {
	host string
	port uint16
}

var alice endpoint = endpoint{"127.0.0.1", 10862}
var bob endpoint = endpoint{"127.0.0.1", 10863}

func TestMassUDP(t *testing.T) {
	NUM_MSGS := 10000
	from, _ := NewManager("udp")
	to, _ := NewManager("udp")
	to.Listen(fmt.Sprintf("%s:%d", alice.host, alice.port))
	receiver := to.GetChannel()

	receivedIDs := make([]int, 0)
	result := make(chan bool)

	go func() {
	recvloop:
		for {
			select {
			case msg, ok := <-receiver:
				if !ok {
					break recvloop
				}
				id, _ := strconv.ParseInt(msg.GetBody(), 10, 32)
				receivedIDs = append(receivedIDs, int(id))
				if len(receivedIDs) >= NUM_MSGS {
					break recvloop
				}
			case <-time.After(time.Second / 10):
				break recvloop
			}
		}

		if !(len(receivedIDs) == NUM_MSGS) {
			t.Errorf("Error - received an unexpected number of messages: %d", len(receivedIDs))
		} else {
			seenSet := make(map[int]bool)
			for _, val := range receivedIDs {
				if _, match := seenSet[val]; match {
					t.Errorf("Duplicate message received: %d", val)
				}
				seenSet[val] = true
			}
			if len(seenSet) != NUM_MSGS {
				t.Errorf("Unexpected number of messages received: %d", len(seenSet))
			}
		}
		result <- true
	}()

	go func() {
		user := "alice"
		uri := base.SipUri{User: &user, Host: "127.0.0.1", Port: nil}
		for ii := 1; ii <= NUM_MSGS; ii++ {
			from.Send(fmt.Sprintf("%s:%d", alice.host, alice.port),
				base.NewRequest(base.ACK, &uri, "SIP/2.0",
					[]base.SipHeader{base.ContentLength(len(fmt.Sprintf("%d", ii)))},
					fmt.Sprintf("%d", ii)))

			if ii%100 == 0 {
				<-time.After(time.Millisecond)
			}
		}
	}()

	_ = <-result
	return
}

func sendAndCheckReceipt(from *Manager, to string,
	receiver chan base.SipMessage,
	msg base.SipMessage, timeout time.Duration) bool {
	from.Send(to, msg)

	select {
	case msgIn, ok := <-receiver:
		if !ok {
			return false
		}

		return msgIn.String() == msg.String()

	case <-time.After(timeout):
		return false
	}
}
<commit_msg>Fixed up transport layer<commit_after>package transport

import (
	"fmt"
	"strconv"
	"testing"
	"time"

	"github.com/stefankopieczek/gossip/base"
)

type endpoint struct {
	host string
	port uint16
}

var alice endpoint = endpoint{"127.0.0.1", 10862}
var bob endpoint = endpoint{"127.0.0.1", 10863}

func TestMassUDP(t *testing.T) {
	NUM_MSGS := 10000
	from, _ := NewManager("udp")
	to, _ := NewManager("udp")
	to.Listen(fmt.Sprintf("%s:%d", alice.host, alice.port))
	receiver := to.GetChannel()

	receivedIDs := make([]int, 0)
	result := make(chan bool)

	go func() {
	recvloop:
		for {
			select {
			case msg, ok := <-receiver:
				if !ok {
					break recvloop
				}
				id, _ := strconv.ParseInt(msg.GetBody(), 10, 32)
				receivedIDs = append(receivedIDs, int(id))
				if len(receivedIDs) >= NUM_MSGS {
					break recvloop
				}
			case <-time.After(time.Second / 10):
				break recvloop
			}
		}

		if !(len(receivedIDs) == NUM_MSGS) {
			t.Errorf("Error - received an unexpected number of messages: %d", len(receivedIDs))
		} else {
			seenSet := make(map[int]bool)
			for _, val := range receivedIDs {
				if _, match := seenSet[val]; match {
					t.Errorf("Duplicate message received: %d", val)
				}
				seenSet[val] = true
			}
			if len(seenSet) != NUM_MSGS {
				t.Errorf("Unexpected number of messages received: %d", len(seenSet))
			}
		}
		result <- true
	}()

	go func() {
		uri := base.SipUri{User: base.String{"alice"}, Host: "127.0.0.1", Port: nil}
		for ii := 1; ii <= NUM_MSGS; ii++ {
			from.Send(fmt.Sprintf("%s:%d", alice.host, alice.port),
				base.NewRequest(base.ACK, &uri, "SIP/2.0",
					[]base.SipHeader{base.ContentLength(len(fmt.Sprintf("%d", ii)))},
					fmt.Sprintf("%d", ii)))

			if ii%100 == 0 {
				<-time.After(time.Millisecond)
			}
		}
	}()

	_ = <-result
	return
}

func sendAndCheckReceipt(from *Manager, to string,
	receiver chan base.SipMessage,
	msg base.SipMessage, timeout time.Duration) bool {
	from.Send(to, msg)

	select {
	case msgIn, ok := <-receiver:
		if !ok {
			return false
		}

		return msgIn.String() == msg.String()

	case <-time.After(timeout):
		return false
	}
}
<|endoftext|>


<commit_before>package server

import (
	//	"fmt"
	"log"
	"net"

	"fmt"
	"github.com/appcelerator/amp/api/rpc/project"
	"github.com/appcelerator/amp/api/rpc/service"
	"github.com/appcelerator/amp/data/elasticsearch"
	"google.golang.org/grpc"
	"gopkg.in/olivere/elastic.v3"
	"os"
)

var (
	es elasticsearch.ElasticSearch
)

func init() {
	// Get elasticsearch url from environment
	elasticSearchURL := os.Getenv("ES_URL")
	if elasticSearchURL == "" {
		elasticSearchURL = elastic.DefaultURL
	}
	fmt.Printf("ES_URL: %v\n", elasticSearchURL)

	// Initialize elastic search
	es = elasticsearch.ElasticSearch{}
	es.Connect(elasticSearchURL)
	es.CreateIndexIfNotExists(esIndex, esType, esMapping)
}

// Start starts the server
func Start(port string) {
	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	project.RegisterProjectServer(s, &projectService{})
	service.RegisterServiceServer(s, &serviceService{})
	s.Serve(lis)
}
<commit_msg>updated imports<commit_after>package server

import (
	//	"fmt"
	"fmt"
	"log"
	"net"
	"os"

	"github.com/appcelerator/amp/api/rpc/project"
	"github.com/appcelerator/amp/api/rpc/service"
	"github.com/appcelerator/amp/data/elasticsearch"
	"google.golang.org/grpc"
	"gopkg.in/olivere/elastic.v3"
)

var (
	es elasticsearch.ElasticSearch
)

func init() {
	// Get elasticsearch url from environment
	elasticSearchURL := os.Getenv("ES_URL")
	if elasticSearchURL == "" {
		elasticSearchURL = elastic.DefaultURL
	}
	fmt.Printf("ES_URL: %v\n", elasticSearchURL)

	// Initialize elastic search
	es = elasticsearch.ElasticSearch{}
	es.Connect(elasticSearchURL)
	es.CreateIndexIfNotExists(esIndex, esType, esMapping)
}

// Start starts the server
func Start(port string) {
	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	project.RegisterProjectServer(s, &projectService{})
	service.RegisterServiceServer(s, &serviceService{})
	s.Serve(lis)
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/api/errors"
	"github.com/docker/docker/api/server/httputils"
	"github.com/docker/docker/api/server/middleware"
	"github.com/docker/docker/api/server/router"
	"github.com/docker/docker/dockerversion"
	"github.com/gorilla/mux"
	"golang.org/x/net/context"
)

// versionMatcher defines a variable matcher to be parsed by the router
// when a request is about to be served.
const versionMatcher = "/v{version:[0-9.]+}"

// Config provides the configuration for the API server
type Config struct {
	Logging     bool
	EnableCors  bool
	CorsHeaders string
	Version     string
	SocketGroup string
	TLSConfig   *tls.Config
}

// Server contains instance details for the server
type Server struct {
	cfg           *Config
	servers       []*HTTPServer
	routers       []router.Router
	routerSwapper *routerSwapper
	middlewares   []middleware.Middleware
}

// New returns a new instance of the server based on the specified configuration.
// It allocates resources which will be needed for ServeAPI(ports, unix-sockets).
func New(cfg *Config) *Server {
	return &Server{
		cfg: cfg,
	}
}

// UseMiddleware appends a new middleware to the request chain.
// This needs to be called before the API routes are configured.
func (s *Server) UseMiddleware(m middleware.Middleware) {
	s.middlewares = append(s.middlewares, m)
}

// Accept sets a listener the server accepts connections into.
func (s *Server) Accept(addr string, listeners ...net.Listener) {
	for _, listener := range listeners {
		httpServer := &HTTPServer{
			srv: &http.Server{
				Addr: addr,
			},
			l: listener,
		}
		s.servers = append(s.servers, httpServer)
	}
}

// Close closes servers and thus stop receiving requests
func (s *Server) Close() {
	for _, srv := range s.servers {
		if err := srv.Close(); err != nil {
			logrus.Error(err)
		}
	}
}

// serveAPI loops through all initialized servers and spawns goroutine
// with Serve method for each. It sets createMux() as Handler also.
func (s *Server) serveAPI() error {
	var chErrors = make(chan error, len(s.servers))
	for _, srv := range s.servers {
		srv.srv.Handler = s.routerSwapper
		go func(srv *HTTPServer) {
			var err error
			logrus.Infof("API listen on %s", srv.l.Addr())
			if err = srv.Serve(); err != nil && strings.Contains(err.Error(), "use of closed network connection") {
				err = nil
			}
			chErrors <- err
		}(srv)
	}

	for i := 0; i < len(s.servers); i++ {
		err := <-chErrors
		if err != nil {
			return err
		}
	}

	return nil
}

// HTTPServer contains an instance of http server and the listener.
// srv *http.Server, contains configuration to create an http server and a mux router with all api end points.
// l   net.Listener, is a TCP or Socket listener that dispatches incoming request to the router.
type HTTPServer struct {
	srv *http.Server
	l   net.Listener
}

// Serve starts listening for inbound requests.
func (s *HTTPServer) Serve() error {
	return s.srv.Serve(s.l)
}

// Close closes the HTTPServer from listening for the inbound requests.
func (s *HTTPServer) Close() error {
	return s.l.Close()
}

func (s *Server) makeHTTPHandler(handler httputils.APIFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// Define the context that we'll pass around to share info
		// like the docker-request-id.
		//
		// The 'context' will be used for global data that should
		// apply to all requests. Data that is specific to the
		// immediate function being called should still be passed
		// as 'args' on the function call.
		ctx := context.WithValue(context.Background(), dockerversion.UAStringKey, r.Header.Get("User-Agent"))
		handlerFunc := s.handlerWithGlobalMiddlewares(handler)

		vars := mux.Vars(r)
		if vars == nil {
			vars = make(map[string]string)
		}

		if err := handlerFunc(ctx, w, r, vars); err != nil {
			logrus.Errorf("Handler for %s %s returned error: %+v", r.Method, r.URL.Path, err)
			httputils.MakeErrorHandler(err)(w, r)
		}
	}
}

// InitRouter initializes the list of routers for the server.
// This method also enables the Go profiler if enableProfiler is true.
func (s *Server) InitRouter(enableProfiler bool, routers ...router.Router) {
	s.routers = append(s.routers, routers...)

	m := s.createMux()
	if enableProfiler {
		profilerSetup(m)
	}
	s.routerSwapper = &routerSwapper{
		router: m,
	}
}

// createMux initializes the main router the server uses.
func (s *Server) createMux() *mux.Router {
	m := mux.NewRouter()

	logrus.Debug("Registering routers")
	for _, apiRouter := range s.routers {
		for _, r := range apiRouter.Routes() {
			f := s.makeHTTPHandler(r.Handler())

			logrus.Debugf("Registering %s, %s", r.Method(), r.Path())
			m.Path(versionMatcher + r.Path()).Methods(r.Method()).Handler(f)
			m.Path(r.Path()).Methods(r.Method()).Handler(f)
		}
	}

	err := errors.NewRequestNotFoundError(fmt.Errorf("page not found"))
	notFoundHandler := httputils.MakeErrorHandler(err)
	m.HandleFunc(versionMatcher+"/{path:.*}", notFoundHandler)
	m.NotFoundHandler = notFoundHandler

	return m
}

// Wait blocks the server goroutine until it exits.
// It sends an error message if there is any error during
// the API execution.
func (s *Server) Wait(waitChan chan error) {
	if err := s.serveAPI(); err != nil {
		logrus.Errorf("ServeAPI error: %v", err)
		waitChan <- err
		return
	}
	waitChan <- nil
}

// DisableProfiler reloads the server mux without adding the profiler routes.
func (s *Server) DisableProfiler() {
	s.routerSwapper.Swap(s.createMux())
}

// EnableProfiler reloads the server mux adding the profiler routes.
func (s *Server) EnableProfiler() {
	m := s.createMux()
	profilerSetup(m)
	s.routerSwapper.Swap(m)
}
<commit_msg>Show stacktrace in daemon logs only if 500 internal error<commit_after>package server

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/api/errors"
	"github.com/docker/docker/api/server/httputils"
	"github.com/docker/docker/api/server/middleware"
	"github.com/docker/docker/api/server/router"
	"github.com/docker/docker/dockerversion"
	"github.com/gorilla/mux"
	"golang.org/x/net/context"
)

// versionMatcher defines a variable matcher to be parsed by the router
// when a request is about to be served.
const versionMatcher = "/v{version:[0-9.]+}"

// Config provides the configuration for the API server
type Config struct {
	Logging     bool
	EnableCors  bool
	CorsHeaders string
	Version     string
	SocketGroup string
	TLSConfig   *tls.Config
}

// Server contains instance details for the server
type Server struct {
	cfg           *Config
	servers       []*HTTPServer
	routers       []router.Router
	routerSwapper *routerSwapper
	middlewares   []middleware.Middleware
}

// New returns a new instance of the server based on the specified configuration.
// It allocates resources which will be needed for ServeAPI(ports, unix-sockets).
func New(cfg *Config) *Server {
	return &Server{
		cfg: cfg,
	}
}

// UseMiddleware appends a new middleware to the request chain.
// This needs to be called before the API routes are configured.
func (s *Server) UseMiddleware(m middleware.Middleware) {
	s.middlewares = append(s.middlewares, m)
}

// Accept sets a listener the server accepts connections into.
func (s *Server) Accept(addr string, listeners ...net.Listener) {
	for _, listener := range listeners {
		httpServer := &HTTPServer{
			srv: &http.Server{
				Addr: addr,
			},
			l: listener,
		}
		s.servers = append(s.servers, httpServer)
	}
}

// Close closes servers and thus stop receiving requests
func (s *Server) Close() {
	for _, srv := range s.servers {
		if err := srv.Close(); err != nil {
			logrus.Error(err)
		}
	}
}

// serveAPI loops through all initialized servers and spawns goroutine
// with Serve method for each. It sets createMux() as Handler also.
func (s *Server) serveAPI() error {
	var chErrors = make(chan error, len(s.servers))
	for _, srv := range s.servers {
		srv.srv.Handler = s.routerSwapper
		go func(srv *HTTPServer) {
			var err error
			logrus.Infof("API listen on %s", srv.l.Addr())
			if err = srv.Serve(); err != nil && strings.Contains(err.Error(), "use of closed network connection") {
				err = nil
			}
			chErrors <- err
		}(srv)
	}

	for i := 0; i < len(s.servers); i++ {
		err := <-chErrors
		if err != nil {
			return err
		}
	}

	return nil
}

// HTTPServer contains an instance of http server and the listener.
// srv *http.Server, contains configuration to create an http server and a mux router with all api end points.
// l   net.Listener, is a TCP or Socket listener that dispatches incoming request to the router.
type HTTPServer struct {
	srv *http.Server
	l   net.Listener
}

// Serve starts listening for inbound requests.
func (s *HTTPServer) Serve() error {
	return s.srv.Serve(s.l)
}

// Close closes the HTTPServer from listening for the inbound requests.
func (s *HTTPServer) Close() error {
	return s.l.Close()
}

func (s *Server) makeHTTPHandler(handler httputils.APIFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// Define the context that we'll pass around to share info
		// like the docker-request-id.
		//
		// The 'context' will be used for global data that should
		// apply to all requests. Data that is specific to the
		// immediate function being called should still be passed
		// as 'args' on the function call.
		ctx := context.WithValue(context.Background(), dockerversion.UAStringKey, r.Header.Get("User-Agent"))
		handlerFunc := s.handlerWithGlobalMiddlewares(handler)

		vars := mux.Vars(r)
		if vars == nil {
			vars = make(map[string]string)
		}

		if err := handlerFunc(ctx, w, r, vars); err != nil {
			statusCode := httputils.GetHTTPErrorStatusCode(err)
			errFormat := "%v"
			if statusCode == http.StatusInternalServerError {
				errFormat = "%+v"
			}
			logrus.Errorf("Handler for %s %s returned error: "+errFormat, r.Method, r.URL.Path, err)
			httputils.MakeErrorHandler(err)(w, r)
		}
	}
}

// InitRouter initializes the list of routers for the server.
// This method also enables the Go profiler if enableProfiler is true.
func (s *Server) InitRouter(enableProfiler bool, routers ...router.Router) {
	s.routers = append(s.routers, routers...)

	m := s.createMux()
	if enableProfiler {
		profilerSetup(m)
	}
	s.routerSwapper = &routerSwapper{
		router: m,
	}
}

// createMux initializes the main router the server uses.
func (s *Server) createMux() *mux.Router {
	m := mux.NewRouter()

	logrus.Debug("Registering routers")
	for _, apiRouter := range s.routers {
		for _, r := range apiRouter.Routes() {
			f := s.makeHTTPHandler(r.Handler())

			logrus.Debugf("Registering %s, %s", r.Method(), r.Path())
			m.Path(versionMatcher + r.Path()).Methods(r.Method()).Handler(f)
			m.Path(r.Path()).Methods(r.Method()).Handler(f)
		}
	}

	err := errors.NewRequestNotFoundError(fmt.Errorf("page not found"))
	notFoundHandler := httputils.MakeErrorHandler(err)
	m.HandleFunc(versionMatcher+"/{path:.*}", notFoundHandler)
	m.NotFoundHandler = notFoundHandler

	return m
}

// Wait blocks the server goroutine until it exits.
// It sends an error message if there is any error during
// the API execution.
func (s *Server) Wait(waitChan chan error) {
	if err := s.serveAPI(); err != nil {
		logrus.Errorf("ServeAPI error: %v", err)
		waitChan <- err
		return
	}
	waitChan <- nil
}

// DisableProfiler reloads the server mux without adding the profiler routes.
func (s *Server) DisableProfiler() {
	s.routerSwapper.Swap(s.createMux())
}

// EnableProfiler reloads the server mux adding the profiler routes.
func (s *Server) EnableProfiler() {
	m := s.createMux()
	profilerSetup(m)
	s.routerSwapper.Swap(m)
}
<|endoftext|>


<commit_before>// Copyright 2013 Bobby Powers. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package xmile_test

import (
	"bufio"
	"bytes"
	"fmt"
	xmile "github.com/bpowers/go-xmile/compat"
	"io/ioutil"
	"os"
	"testing"
	"text/template"
)

const dotTmpl = `
digraph model {
{{range $.Clouds}}
{{.}} [shape=none,label=""]{{end}}

{{range $.Stocks}}
{{.Name}} [shape=box]{{end}}

{{range $.Flows}}
{{.Name}} [shape=circle]{{end}}

{{range $.Auxs}}
{{.Name}} [shape=circle]{{end}}

{{range $.Stocks}}{{with $s := .}}{{range $.OutsC}}
{{$s.Name}} -> {{.}}{{end}}{{end}}{{end}}

{{range $.Auxs}}{{with $s := .}}{{range $.OutsC}}
{{$s.Name}} -> {{.}}{{end}}{{end}}{{end}}
}
`

type DotData struct {
	Clouds []string
	Stocks []VInfo
	Flows  []VInfo
	Auxs   []VInfo
}

// VInfo stores information about a var for use by dot.
//
// stock: only connectors out.
// flow: all 3, connector outs & flow ins and outs.
// aux: only outsc.
type VInfo struct {
	Name  string
	OutsC []string
	OutsF []string
	Ins   []string // only flows have ins
}

func writeDot(f *xmile.File) error {
	w := bufio.NewWriter(os.Stderr)
	defer w.Flush()

	var data DotData

	var buf bytes.Buffer
	tmpl := template.New("model.dot")
	if _, err := tmpl.Parse(dotTmpl); err != nil {
		return fmt.Errorf("Parse(dotTmpl): %s", err)
	}
	if err := tmpl.Execute(&buf, &data); err != nil {
		return fmt.Errorf("Execute: %s", err)
	}

	w.Write(buf.Bytes())
	w.Write([]byte("\n"))

	return nil
}

/*
func TestRead(t *testing.T) {
	contents, err := ioutil.ReadFile("../models/pred_prey.stmx")
	if err != nil {
		t.Fatalf("ioutil.ReadFile: %s", err)
	}

	f, err := xmile.ReadFile(contents)
	if err != nil {
		t.Fatalf("xmile.ReadFile: %s", err)
	}

	output, err := xml.MarshalIndent(f, "", "    ")
	if err != nil {
		t.Fatalf("xml.MarshalIndent: %s", err)
	}

	os.Stderr.Write([]byte(xmile.XMLDeclaration + "\n"))
	os.Stderr.Write(output)
	os.Stderr.Write([]byte("\n"))
}
*/

func TestDot(t *testing.T) {
	contents, err := ioutil.ReadFile("../models/pred_prey.stmx")
	if err != nil {
		t.Fatalf("ioutil.ReadFile: %s", err)
	}

	f, err := xmile.ReadFile(contents)
	if err != nil {
		t.Fatalf("xmile.ReadFile: %s", err)
	}

	f.Models[0].Interface = xmile.View{}
	if err := writeDot(f); err != nil {
		t.Fatalf("writeDot: %s", err)
	}
}
<commit_msg>normalize equation names<commit_after>// Copyright 2013 Bobby Powers. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package xmile_test

import (
	"bufio"
	"bytes"
	"fmt"
	xmile "github.com/bpowers/go-xmile/compat"
	"io/ioutil"
	"os"
	"strings"
	"testing"
	"text/template"
)

const dotTmpl = `
digraph model {
{{range $.Clouds}}
{{.}} [shape=none,label=""]{{end}}

{{range $.Stocks}}
{{.Name}} [shape=box]{{end}}

{{range $.Flows}}
{{.Name}} [shape=circle]{{end}}

{{range $.Auxs}}
{{.Name}} [shape=circle]{{end}}

{{range $.Stocks}}{{with $s := .}}{{range $.OutsC}}
{{$s.Name}} -> {{.}}{{end}}{{end}}{{end}}

{{range $.Auxs}}{{with $s := .}}{{range $.OutsC}}
{{$s.Name}} -> {{.}}{{end}}{{end}}{{end}}
}
`

type DotData struct {
	Clouds []string
	Stocks []VInfo
	Flows  []VInfo
	Auxs   []VInfo
}

// VInfo stores information about a var for use by dot.
//
// stock: only connectors out.
// flow: all 3, connector outs & flow ins and outs.
// aux: only outsc.
type VInfo struct {
	Name  string
	OutsC []string
	OutsF []string
	Ins   []string // only flows have ins
}

func normalizeName(n string) string {
	n = strings.Replace(n, `\n`, "_", -1)
	n = strings.ToLower(n)
	return n
}

func normalizeNames(f *xmile.File) {
	for _, m := range f.Models {
		for _, v := range m.Variables {
			v.Name = normalizeName(v.Name)
		}
	}
}

func writeDot(f *xmile.File) error {
	normalizeNames(f)

	var data DotData

	var buf bytes.Buffer
	tmpl := template.New("model.dot")
	if _, err := tmpl.Parse(dotTmpl); err != nil {
		return fmt.Errorf("Parse(dotTmpl): %s", err)
	}
	if err := tmpl.Execute(&buf, &data); err != nil {
		return fmt.Errorf("Execute: %s", err)
	}

	w := bufio.NewWriter(os.Stderr)
	defer w.Flush()
	w.Write(buf.Bytes())
	w.Write([]byte("\n"))

	return nil
}

/*
func TestRead(t *testing.T) {
	contents, err := ioutil.ReadFile("../models/pred_prey.stmx")
	if err != nil {
		t.Fatalf("ioutil.ReadFile: %s", err)
	}

	f, err := xmile.ReadFile(contents)
	if err != nil {
		t.Fatalf("xmile.ReadFile: %s", err)
	}

	output, err := xml.MarshalIndent(f, "", "    ")
	if err != nil {
		t.Fatalf("xml.MarshalIndent: %s", err)
	}

	os.Stderr.Write([]byte(xmile.XMLDeclaration + "\n"))
	os.Stderr.Write(output)
	os.Stderr.Write([]byte("\n"))
}
*/

func TestDot(t *testing.T) {
	contents, err := ioutil.ReadFile("../models/pred_prey.stmx")
	if err != nil {
		t.Fatalf("ioutil.ReadFile: %s", err)
	}

	f, err := xmile.ReadFile(contents)
	if err != nil {
		t.Fatalf("xmile.ReadFile: %s", err)
	}

	f.Models[0].Interface = xmile.View{}
	if err := writeDot(f); err != nil {
		t.Fatalf("writeDot: %s", err)
	}
}
<|endoftext|>


<commit_before>package api_ws

import (
	"fmt"
	"sync"
	"time"

	"github.com/deze333/vroom/util"
	"github.com/gorilla/websocket"
)

//------------------------------------------------------------
// WebSocket messages async processor
//------------------------------------------------------------

const (
	CHAN_MESSAGE_MAX           = 1024
	CHAN_MESSAGE_BROADCAST_MAX = 100
	CHAN_MESSAGE_FAILED_MAX    = 100
	LOST_MESSAGE_MAX_LIFE      = 30 * time.Second
)

var (
	_chanIn          = make(chan *Message, CHAN_MESSAGE_MAX)
	_chanOut         = make(chan *Message, CHAN_MESSAGE_MAX)
	_chanBroadcast   = make(chan *Message, CHAN_MESSAGE_BROADCAST_MAX)
	_chanFailed      = make(chan *Message, CHAN_MESSAGE_FAILED_MAX)
	_failed          = map[string][]*Message{}
	_procFailedMutex = &sync.Mutex{}
)

// Initializes async processors.
func init() {
	go procChanIn()
	go procChanOut()
	go procChanBroadcast()
	go procChanFailed()
	go procFailedCleanup()
}

// Processes incoming channel messages. One message at a time.
func procChanIn() {
	var conn *Conn
	for msg := range _chanIn {
		msg.isProcessed = false
		//fmt.Printf("chan IN ---> %v = %v\n", msg.agentId, string(msg.req))

		// Attempt to get active connection
		if conn = GetConn(msg.isAuthd, msg.agentId); conn == nil {
			_chanFailed <- msg
			continue
		}

		// Process message
		// Must be done via goroutine to prevent channel block
		// due to processing delay
		go processAndRespond(conn, msg, _chanOut)
	}
}

// Processes outgoing channel messages. One message at a time.
func procChanOut() {
	var conn *Conn
	for msg := range _chanOut {
		//fmt.Printf("chan OUT <--- %v\n", string(msg.res))

		// Attempt to get active connection
		if conn = GetConn(msg.isAuthd, msg.agentId); conn == nil {
			_chanFailed <- msg
			continue
		}

		// Write response
		err := conn.conn.WriteMessage(websocket.TextMessage, msg.res)
		if err != nil {
			// Close connection on write error
			_chanFailed <- msg
			DeregisterConn(conn)
			conn.conn.Close()
			continue
		}
	}
}

// Sends broadcast messages to all active connections.
func procChanBroadcast() {
	for msg := range _chanBroadcast {
		//fmt.Printf("chan BROADCAST <--- %v\n", string(msg.res))

		// Apply broadcaster function to each connection
		if msg.isAuthd {
			applyToAuthd(
				func(conn *Conn) {
					broadcaster(conn, msg)
				})
		} else {
			applyToPublic(
				func(conn *Conn) {
					broadcaster(conn, msg)
				})
		}
	}
}

// Broadcasts message to single connection.
// Closes connection on error.
func broadcaster(conn *Conn, msg *Message) {

	// Catch panic
	defer func() {
		if err := recover(); err != nil {

			// Report panic: err, url, params, stack
			_onPanic(
				fmt.Sprintf("Error writing to WebSocket: %v", err),
				"Stack", util.Stack())
		}
	}()

	err := conn.conn.WriteMessage(websocket.TextMessage, msg.res)
	if err != nil {
		// Close connection on write error
		_chanFailed <- msg
		DeregisterConn(conn)
		conn.conn.Close()
	}
}

// Listens on failed channel and adds messages to failed.
func procChanFailed() {
	for msg := range _chanFailed {
		_procFailedMutex.Lock()

		t := time.Now()
		msg.failTime = &t
		if agentFails, ok := _failed[msg.agentId]; ok {
			_failed[msg.agentId] = append(agentFails, msg)
		} else {
			_failed[msg.agentId] = []*Message{msg}
		}

		_procFailedMutex.Unlock()
	}
}

// Periodically cleans up expired failed messages.
func procFailedCleanup() {

	ticker := time.NewTicker(15 * time.Second)

	for {
		<-ticker.C
		_procFailedMutex.Lock()

		now := time.Now()

		for agentId, agentFails := range _failed {

			if anyExpiredFails(agentFails, now) {
				unexp := getUnexpiredFails(agentFails, now)
				if len(unexp) > 0 {
					_failed[agentId] = unexp
				} else {
					delete(_failed, agentId)
				}
			}
		}

		_procFailedMutex.Unlock()
	}
}

// Has at least one message expired?
func anyExpiredFails(msgs []*Message, t time.Time) bool {
	for _, msg := range msgs {
		if msg.failTime == nil {
			return true // safe guard, nil not allowed
		}
		if msg.failTime.Add(LOST_MESSAGE_MAX_LIFE).After(t) {
			return true
		}
	}
	return false
}

// Filters out expired messages.
func getUnexpiredFails(msgs []*Message, t time.Time) (unexp []*Message) {
	for _, msg := range msgs {
		if msg.failTime == nil {
			continue // safe guard, nil not allowed
		}
		if msg.failTime.Add(LOST_MESSAGE_MAX_LIFE).After(t) {
			continue
		}
		unexp = append(unexp, msg)
	}
	return
}

// Pushes specific agent's failed messages back to channel for processing.
func retryFailedMessages(agentId string) {

	_procFailedMutex.Lock()

	if agentFails, ok := _failed[agentId]; ok {
		delete(_failed, agentId)
		for _, msg := range agentFails {
			msg.failTime = nil
			if msg.isProcessed {
				_chanOut <- msg
			} else {
				_chanIn <- msg
			}
		}
	}

	_procFailedMutex.Unlock()
}
<commit_msg>catch error at websocket write<commit_after>package api_ws

import (
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

//------------------------------------------------------------
// WebSocket messages async processor
//------------------------------------------------------------

const (
	CHAN_MESSAGE_MAX           = 1024
	CHAN_MESSAGE_BROADCAST_MAX = 100
	CHAN_MESSAGE_FAILED_MAX    = 100
	LOST_MESSAGE_MAX_LIFE      = 30 * time.Second
)

var (
	_chanIn          = make(chan *Message, CHAN_MESSAGE_MAX)
	_chanOut         = make(chan *Message, CHAN_MESSAGE_MAX)
	_chanBroadcast   = make(chan *Message, CHAN_MESSAGE_BROADCAST_MAX)
	_chanFailed      = make(chan *Message, CHAN_MESSAGE_FAILED_MAX)
	_failed          = map[string][]*Message{}
	_procFailedMutex = &sync.Mutex{}
)

// Initializes async processors.
func init() {
	go procChanIn()
	go procChanOut()
	go procChanBroadcast()
	go procChanFailed()
	go procFailedCleanup()
}

// Processes incoming channel messages. One message at a time.
func procChanIn() {
	var conn *Conn
	for msg := range _chanIn {
		msg.isProcessed = false
		//fmt.Printf("chan IN ---> %v = %v\n", msg.agentId, string(msg.req))

		// Attempt to get active connection
		if conn = GetConn(msg.isAuthd, msg.agentId); conn == nil {
			_chanFailed <- msg
			continue
		}

		// Process message
		// Must be done via goroutine to prevent channel block
		// due to processing delay
		go processAndRespond(conn, msg, _chanOut)
	}
}

// Processes outgoing channel messages. One message at a time.
func procChanOut() {
	var conn *Conn
	for msg := range _chanOut {
		//fmt.Printf("chan OUT <--- %v\n", string(msg.res))

		// Attempt to get active connection
		if conn = GetConn(msg.isAuthd, msg.agentId); conn == nil {
			_chanFailed <- msg
			continue
		}

		// Write response
		err := conn.conn.WriteMessage(websocket.TextMessage, msg.res)
		if err != nil {
			// Close connection on write error
			_chanFailed <- msg
			DeregisterConn(conn)
			conn.conn.Close()
			continue
		}
	}
}

// Sends broadcast messages to all active connections.
func procChanBroadcast() {
	for msg := range _chanBroadcast {
		//fmt.Printf("chan BROADCAST <--- %v\n", string(msg.res))

		// Apply broadcaster function to each connection
		if msg.isAuthd {
			applyToAuthd(
				func(conn *Conn) {
					broadcaster(conn, msg)
				})
		} else {
			applyToPublic(
				func(conn *Conn) {
					broadcaster(conn, msg)
				})
		}
	}
}

// Broadcasts message to single connection.
// Closes connection on error.
func broadcaster(conn *Conn, msg *Message) {

	// Catch panic
	defer func() {
		if err := recover(); err != nil {

			// Do nothing
			// Report panic: err, url, params, stack
			/*
				_onPanic(
					fmt.Sprintf("Error writing to WebSocket: %v", err),
					"Stack", util.Stack())
			*/
		}
	}()

	err := conn.conn.WriteMessage(websocket.TextMessage, msg.res)
	if err != nil {
		// Close connection on write error
		_chanFailed <- msg
		DeregisterConn(conn)
		conn.conn.Close()
	}
}

// Listens on failed channel and adds messages to failed.
func procChanFailed() {
	for msg := range _chanFailed {
		_procFailedMutex.Lock()

		t := time.Now()
		msg.failTime = &t
		if agentFails, ok := _failed[msg.agentId]; ok {
			_failed[msg.agentId] = append(agentFails, msg)
		} else {
			_failed[msg.agentId] = []*Message{msg}
		}

		_procFailedMutex.Unlock()
	}
}

// Periodically cleans up expired failed messages.
func procFailedCleanup() {

	ticker := time.NewTicker(15 * time.Second)

	for {
		<-ticker.C
		_procFailedMutex.Lock()

		now := time.Now()

		for agentId, agentFails := range _failed {

			if anyExpiredFails(agentFails, now) {
				unexp := getUnexpiredFails(agentFails, now)
				if len(unexp) > 0 {
					_failed[agentId] = unexp
				} else {
					delete(_failed, agentId)
				}
			}
		}

		_procFailedMutex.Unlock()
	}
}

// Has at least one message expired?
func anyExpiredFails(msgs []*Message, t time.Time) bool {
	for _, msg := range msgs {
		if msg.failTime == nil {
			return true // safe guard, nil not allowed
		}
		if msg.failTime.Add(LOST_MESSAGE_MAX_LIFE).After(t) {
			return true
		}
	}
	return false
}

// Filters out expired messages.
func getUnexpiredFails(msgs []*Message, t time.Time) (unexp []*Message) {
	for _, msg := range msgs {
		if msg.failTime == nil {
			continue // safe guard, nil not allowed
		}
		if msg.failTime.Add(LOST_MESSAGE_MAX_LIFE).After(t) {
			continue
		}
		unexp = append(unexp, msg)
	}
	return
}

// Pushes specific agent's failed messages back to channel for processing.
func retryFailedMessages(agentId string) {

	_procFailedMutex.Lock()

	if agentFails, ok := _failed[agentId]; ok {
		delete(_failed, agentId)
		for _, msg := range agentFails {
			msg.failTime = nil
			if msg.isProcessed {
				_chanOut <- msg
			} else {
				_chanIn <- msg
			}
		}
	}

	_procFailedMutex.Unlock()
}
<|endoftext|>