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

package rpc

import (
	"bufio"
	"encoding/gob"
	"errors"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
)

// ServerError represents an error that has been returned from
// the remote side of the RPC connection.
type ServerError string

func (e ServerError) Error() string {
	return string(e)
}

var ErrShutdown = errors.New("connection is shut down")

// Call represents an active RPC.
type Call struct {
	ServiceMethod string      // The name of the service and method to call.
	Args          interface{} // The argument to the function (*struct).
	Reply         interface{} // The reply from the function (*struct).
	Error         error       // After completion, the error status.
	Done          chan *Call  // Strobes when call is complete; value is the error status.
	seq           uint64
}

// Client represents an RPC Client.
// There may be multiple outstanding Calls associated
// with a single Client.
type Client struct {
	mutex    sync.Mutex // protects pending, seq, request
	sending  sync.Mutex
	request  Request
	seq      uint64
	codec    ClientCodec
	pending  map[uint64]*Call
	closing  bool
	shutdown bool
}

// A ClientCodec implements writing of RPC requests and
// reading of RPC responses for the client side of an RPC session.
// The client calls WriteRequest to write a request to the connection
// and calls ReadResponseHeader and ReadResponseBody in pairs
// to read responses.  The client calls Close when finished with the
// connection. ReadResponseBody may be called with a nil
// argument to force the body of the response to be read and then
// discarded.
type ClientCodec interface {
	WriteRequest(*Request, interface{}) error
	ReadResponseHeader(*Response) error
	ReadResponseBody(interface{}) error

	Close() error
}

func (client *Client) send(c *Call) {
	// Register this call.
	client.mutex.Lock()
	if client.shutdown {
		c.Error = ErrShutdown
		client.mutex.Unlock()
		c.done()
		return
	}
	c.seq = client.seq
	client.seq++
	client.pending[c.seq] = c
	client.mutex.Unlock()

	// Encode and send the request.
	client.sending.Lock()
	defer client.sending.Unlock()
	client.request.Seq = c.seq
	client.request.ServiceMethod = c.ServiceMethod
	if err := client.codec.WriteRequest(&client.request, c.Args); err != nil {
		c.Error = err
		c.done()
	}
}

func (client *Client) input() {
	var err error
	var response Response
	for err == nil {
		response = Response{}
		err = client.codec.ReadResponseHeader(&response)
		if err != nil {
			if err == io.EOF && !client.closing {
				err = io.ErrUnexpectedEOF
			}
			break
		}
		seq := response.Seq
		client.mutex.Lock()
		c := client.pending[seq]
		delete(client.pending, seq)
		client.mutex.Unlock()

		if response.Error == "" {
			err = client.codec.ReadResponseBody(c.Reply)
			if err != nil {
				c.Error = errors.New("reading body " + err.Error())
			}
		} else {
			// We've got an error response. Give this to the request;
			// any subsequent requests will get the ReadResponseBody
			// error if there is one.
			c.Error = ServerError(response.Error)
			err = client.codec.ReadResponseBody(nil)
			if err != nil {
				err = errors.New("reading error body: " + err.Error())
			}
		}
		c.done()
	}
	// Terminate pending calls.
	client.mutex.Lock()
	client.shutdown = true
	for _, call := range client.pending {
		call.Error = err
		call.done()
	}
	client.mutex.Unlock()
	if err != io.EOF || !client.closing {
		log.Println("rpc: client protocol error:", err)
	}
}

func (call *Call) done() {
	select {
	case call.Done <- call:
		// ok
	default:
		// We don't want to block here.  It is the caller's responsibility to make
		// sure the channel has enough buffer space. See comment in Go().
	}
}

// NewClient returns a new Client to handle requests to the
// set of services at the other end of the connection.
// It adds a buffer to the write side of the connection so
// the header and payload are sent as a unit.
func NewClient(conn io.ReadWriteCloser) *Client {
	encBuf := bufio.NewWriter(conn)
	client := &gobClientCodec{conn, gob.NewDecoder(conn), gob.NewEncoder(encBuf), encBuf}
	return NewClientWithCodec(client)
}

// NewClientWithCodec is like NewClient but uses the specified
// codec to encode requests and decode responses.
func NewClientWithCodec(codec ClientCodec) *Client {
	client := &Client{
		codec:   codec,
		pending: make(map[uint64]*Call),
	}
	go client.input()
	return client
}

type gobClientCodec struct {
	rwc    io.ReadWriteCloser
	dec    *gob.Decoder
	enc    *gob.Encoder
	encBuf *bufio.Writer
}

func (c *gobClientCodec) WriteRequest(r *Request, body interface{}) (err error) {
	if err = c.enc.Encode(r); err != nil {
		return
	}
	if err = c.enc.Encode(body); err != nil {
		return
	}
	return c.encBuf.Flush()
}

func (c *gobClientCodec) ReadResponseHeader(r *Response) error {
	return c.dec.Decode(r)
}

func (c *gobClientCodec) ReadResponseBody(body interface{}) error {
	return c.dec.Decode(body)
}

func (c *gobClientCodec) Close() error {
	return c.rwc.Close()
}

// DialHTTP connects to an HTTP RPC server at the specified network address
// listening on the default HTTP RPC path.
func DialHTTP(network, address string) (*Client, error) {
	return DialHTTPPath(network, address, DefaultRPCPath)
}

// DialHTTPPath connects to an HTTP RPC server
// at the specified network address and path.
func DialHTTPPath(network, address, path string) (*Client, error) {
	var err error
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	io.WriteString(conn, "CONNECT "+path+" HTTP/1.0\n\n")

	// Require successful HTTP response
	// before switching to RPC protocol.
	resp, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: "CONNECT"})
	if err == nil && resp.Status == connected {
		return NewClient(conn), nil
	}
	if err == nil {
		err = errors.New("unexpected HTTP response: " + resp.Status)
	}
	conn.Close()
	return nil, &net.OpError{"dial-http", network + " " + address, nil, err}
}

// Dial connects to an RPC server at the specified network address.
func Dial(network, address string) (*Client, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return NewClient(conn), nil
}

func (client *Client) Close() error {
	client.mutex.Lock()
	if client.shutdown || client.closing {
		client.mutex.Unlock()
		return ErrShutdown
	}
	client.closing = true
	client.mutex.Unlock()
	return client.codec.Close()
}

// Go invokes the function asynchronously.  It returns the Call structure representing
// the invocation.  The done channel will signal when the call is complete by returning
// the same Call object.  If done is nil, Go will allocate a new channel.
// If non-nil, done must be buffered or Go will deliberately crash.
func (client *Client) Go(serviceMethod string, args interface{}, reply interface{}, done chan *Call) *Call {
	call := new(Call)
	call.ServiceMethod = serviceMethod
	call.Args = args
	call.Reply = reply
	if done == nil {
		done = make(chan *Call, 10) // buffered.
	} else {
		// If caller passes done != nil, it must arrange that
		// done has enough buffer for the number of simultaneous
		// RPCs that will be using that channel.  If the channel
		// is totally unbuffered, it's best not to run at all.
		if cap(done) == 0 {
			log.Panic("rpc: done channel is unbuffered")
		}
	}
	call.Done = done
	if client.shutdown {
		call.Error = ErrShutdown
		call.done()
		return call
	}
	client.send(call)
	return call
}

// Call invokes the named function, waits for it to complete, and returns its error status.
func (client *Client) Call(serviceMethod string, args interface{}, reply interface{}) error {
	if client.shutdown {
		return ErrShutdown
	}
	call := <-client.Go(serviceMethod, args, reply, make(chan *Call, 1)).Done
	return call.Error
}
<commit_msg>net/rpc: log Call reply discard It means serious user error that can lead to hard to debug issues under load, log entry will not harm.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"bufio"
	"encoding/gob"
	"errors"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
)

// ServerError represents an error that has been returned from
// the remote side of the RPC connection.
type ServerError string

func (e ServerError) Error() string {
	return string(e)
}

var ErrShutdown = errors.New("connection is shut down")

// Call represents an active RPC.
type Call struct {
	ServiceMethod string      // The name of the service and method to call.
	Args          interface{} // The argument to the function (*struct).
	Reply         interface{} // The reply from the function (*struct).
	Error         error       // After completion, the error status.
	Done          chan *Call  // Strobes when call is complete; value is the error status.
	seq           uint64
}

// Client represents an RPC Client.
// There may be multiple outstanding Calls associated
// with a single Client.
type Client struct {
	mutex    sync.Mutex // protects pending, seq, request
	sending  sync.Mutex
	request  Request
	seq      uint64
	codec    ClientCodec
	pending  map[uint64]*Call
	closing  bool
	shutdown bool
}

// A ClientCodec implements writing of RPC requests and
// reading of RPC responses for the client side of an RPC session.
// The client calls WriteRequest to write a request to the connection
// and calls ReadResponseHeader and ReadResponseBody in pairs
// to read responses.  The client calls Close when finished with the
// connection. ReadResponseBody may be called with a nil
// argument to force the body of the response to be read and then
// discarded.
type ClientCodec interface {
	WriteRequest(*Request, interface{}) error
	ReadResponseHeader(*Response) error
	ReadResponseBody(interface{}) error

	Close() error
}

func (client *Client) send(c *Call) {
	// Register this call.
	client.mutex.Lock()
	if client.shutdown {
		c.Error = ErrShutdown
		client.mutex.Unlock()
		c.done()
		return
	}
	c.seq = client.seq
	client.seq++
	client.pending[c.seq] = c
	client.mutex.Unlock()

	// Encode and send the request.
	client.sending.Lock()
	defer client.sending.Unlock()
	client.request.Seq = c.seq
	client.request.ServiceMethod = c.ServiceMethod
	if err := client.codec.WriteRequest(&client.request, c.Args); err != nil {
		c.Error = err
		c.done()
	}
}

func (client *Client) input() {
	var err error
	var response Response
	for err == nil {
		response = Response{}
		err = client.codec.ReadResponseHeader(&response)
		if err != nil {
			if err == io.EOF && !client.closing {
				err = io.ErrUnexpectedEOF
			}
			break
		}
		seq := response.Seq
		client.mutex.Lock()
		c := client.pending[seq]
		delete(client.pending, seq)
		client.mutex.Unlock()

		if response.Error == "" {
			err = client.codec.ReadResponseBody(c.Reply)
			if err != nil {
				c.Error = errors.New("reading body " + err.Error())
			}
		} else {
			// We've got an error response. Give this to the request;
			// any subsequent requests will get the ReadResponseBody
			// error if there is one.
			c.Error = ServerError(response.Error)
			err = client.codec.ReadResponseBody(nil)
			if err != nil {
				err = errors.New("reading error body: " + err.Error())
			}
		}
		c.done()
	}
	// Terminate pending calls.
	client.mutex.Lock()
	client.shutdown = true
	for _, call := range client.pending {
		call.Error = err
		call.done()
	}
	client.mutex.Unlock()
	if err != io.EOF || !client.closing {
		log.Println("rpc: client protocol error:", err)
	}
}

func (call *Call) done() {
	select {
	case call.Done <- call:
		// ok
	default:
		// We don't want to block here.  It is the caller's responsibility to make
		// sure the channel has enough buffer space. See comment in Go().
		log.Println("rpc: discarding Call reply due to insufficient Done chan capacity")
	}
}

// NewClient returns a new Client to handle requests to the
// set of services at the other end of the connection.
// It adds a buffer to the write side of the connection so
// the header and payload are sent as a unit.
func NewClient(conn io.ReadWriteCloser) *Client {
	encBuf := bufio.NewWriter(conn)
	client := &gobClientCodec{conn, gob.NewDecoder(conn), gob.NewEncoder(encBuf), encBuf}
	return NewClientWithCodec(client)
}

// NewClientWithCodec is like NewClient but uses the specified
// codec to encode requests and decode responses.
func NewClientWithCodec(codec ClientCodec) *Client {
	client := &Client{
		codec:   codec,
		pending: make(map[uint64]*Call),
	}
	go client.input()
	return client
}

type gobClientCodec struct {
	rwc    io.ReadWriteCloser
	dec    *gob.Decoder
	enc    *gob.Encoder
	encBuf *bufio.Writer
}

func (c *gobClientCodec) WriteRequest(r *Request, body interface{}) (err error) {
	if err = c.enc.Encode(r); err != nil {
		return
	}
	if err = c.enc.Encode(body); err != nil {
		return
	}
	return c.encBuf.Flush()
}

func (c *gobClientCodec) ReadResponseHeader(r *Response) error {
	return c.dec.Decode(r)
}

func (c *gobClientCodec) ReadResponseBody(body interface{}) error {
	return c.dec.Decode(body)
}

func (c *gobClientCodec) Close() error {
	return c.rwc.Close()
}

// DialHTTP connects to an HTTP RPC server at the specified network address
// listening on the default HTTP RPC path.
func DialHTTP(network, address string) (*Client, error) {
	return DialHTTPPath(network, address, DefaultRPCPath)
}

// DialHTTPPath connects to an HTTP RPC server
// at the specified network address and path.
func DialHTTPPath(network, address, path string) (*Client, error) {
	var err error
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	io.WriteString(conn, "CONNECT "+path+" HTTP/1.0\n\n")

	// Require successful HTTP response
	// before switching to RPC protocol.
	resp, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: "CONNECT"})
	if err == nil && resp.Status == connected {
		return NewClient(conn), nil
	}
	if err == nil {
		err = errors.New("unexpected HTTP response: " + resp.Status)
	}
	conn.Close()
	return nil, &net.OpError{"dial-http", network + " " + address, nil, err}
}

// Dial connects to an RPC server at the specified network address.
func Dial(network, address string) (*Client, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return NewClient(conn), nil
}

func (client *Client) Close() error {
	client.mutex.Lock()
	if client.shutdown || client.closing {
		client.mutex.Unlock()
		return ErrShutdown
	}
	client.closing = true
	client.mutex.Unlock()
	return client.codec.Close()
}

// Go invokes the function asynchronously.  It returns the Call structure representing
// the invocation.  The done channel will signal when the call is complete by returning
// the same Call object.  If done is nil, Go will allocate a new channel.
// If non-nil, done must be buffered or Go will deliberately crash.
func (client *Client) Go(serviceMethod string, args interface{}, reply interface{}, done chan *Call) *Call {
	call := new(Call)
	call.ServiceMethod = serviceMethod
	call.Args = args
	call.Reply = reply
	if done == nil {
		done = make(chan *Call, 10) // buffered.
	} else {
		// If caller passes done != nil, it must arrange that
		// done has enough buffer for the number of simultaneous
		// RPCs that will be using that channel.  If the channel
		// is totally unbuffered, it's best not to run at all.
		if cap(done) == 0 {
			log.Panic("rpc: done channel is unbuffered")
		}
	}
	call.Done = done
	if client.shutdown {
		call.Error = ErrShutdown
		call.done()
		return call
	}
	client.send(call)
	return call
}

// Call invokes the named function, waits for it to complete, and returns its error status.
func (client *Client) Call(serviceMethod string, args interface{}, reply interface{}) error {
	if client.shutdown {
		return ErrShutdown
	}
	call := <-client.Go(serviceMethod, args, reply, make(chan *Call, 1)).Done
	return call.Error
}
<|endoftext|>


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

// +build darwin freebsd linux netbsd openbsd windows

package net

import (
	"syscall"
	"time"
)

// A sockaddr represents a TCP, UDP, IP or Unix network endpoint
// address that can be converted into a syscall.Sockaddr.
type sockaddr interface {
	Addr
	family() int
	isWildcard() bool
	sockaddr(family int) (syscall.Sockaddr, error)
	toAddr() sockaddr
}

// Generic POSIX socket creation.
func socket(net string, f, t, p int, ipv6only bool, ulsa, ursa syscall.Sockaddr, deadline time.Time, toAddr func(syscall.Sockaddr) Addr) (fd *netFD, err error) {
	s, err := sysSocket(f, t, p)
	if err != nil {
		return nil, err
	}

	if err = setDefaultSockopts(s, f, t, ipv6only); err != nil {
		closesocket(s)
		return nil, err
	}

	// This socket is used by a listener.
	if ulsa != nil && ursa == nil {
		// We provide a socket that listens to a wildcard
		// address with reusable UDP port when the given ulsa
		// is an appropriate UDP multicast address prefix.
		// This makes it possible for a single UDP listener
		// to join multiple different group addresses, for
		// multiple UDP listeners that listen on the same UDP
		// port to join the same group address.
		if ulsa, err = listenerSockaddr(s, f, ulsa, toAddr); err != nil {
			closesocket(s)
			return nil, err
		}
	}

	if ulsa != nil {
		if err = syscall.Bind(s, ulsa); err != nil {
			closesocket(s)
			return nil, err
		}
	}

	if fd, err = newFD(s, f, t, net); err != nil {
		closesocket(s)
		return nil, err
	}

	// This socket is used by a dialer.
	if ursa != nil {
		if !deadline.IsZero() {
			setWriteDeadline(fd, deadline)
		}
		if err = fd.connect(ulsa, ursa); err != nil {
			fd.Close()
			return nil, err
		}
		fd.isConnected = true
		if !deadline.IsZero() {
			setWriteDeadline(fd, noDeadline)
		}
	}

	lsa, _ := syscall.Getsockname(s)
	laddr := toAddr(lsa)
	rsa, _ := syscall.Getpeername(s)
	if rsa == nil {
		rsa = ursa
	}
	raddr := toAddr(rsa)
	fd.setAddr(laddr, raddr)
	return fd, nil
}
<commit_msg>net: document sockaddr interface<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin freebsd linux netbsd openbsd windows

package net

import (
	"syscall"
	"time"
)

// A sockaddr represents a TCP, UDP, IP or Unix network endpoint
// address that can be converted into a syscall.Sockaddr.
type sockaddr interface {
	Addr

	// family returns the platform-dependent address family
	// identifier.
	family() int

	// isWildcard reports whether the address is a wildcard
	// address.
	isWildcard() bool

	// sockaddr returns the address converted into a syscall
	// sockaddr type that implements syscall.Sockaddr
	// interface. It returns a nil interface when the address is
	// nil.
	sockaddr(family int) (syscall.Sockaddr, error)

	// toAddr returns the address represented in sockaddr
	// interface. It returns a nil interface when the address is
	// nil.
	toAddr() sockaddr
}

// Generic POSIX socket creation.
func socket(net string, f, t, p int, ipv6only bool, ulsa, ursa syscall.Sockaddr, deadline time.Time, toAddr func(syscall.Sockaddr) Addr) (fd *netFD, err error) {
	s, err := sysSocket(f, t, p)
	if err != nil {
		return nil, err
	}

	if err = setDefaultSockopts(s, f, t, ipv6only); err != nil {
		closesocket(s)
		return nil, err
	}

	// This socket is used by a listener.
	if ulsa != nil && ursa == nil {
		// We provide a socket that listens to a wildcard
		// address with reusable UDP port when the given ulsa
		// is an appropriate UDP multicast address prefix.
		// This makes it possible for a single UDP listener
		// to join multiple different group addresses, for
		// multiple UDP listeners that listen on the same UDP
		// port to join the same group address.
		if ulsa, err = listenerSockaddr(s, f, ulsa, toAddr); err != nil {
			closesocket(s)
			return nil, err
		}
	}

	if ulsa != nil {
		if err = syscall.Bind(s, ulsa); err != nil {
			closesocket(s)
			return nil, err
		}
	}

	if fd, err = newFD(s, f, t, net); err != nil {
		closesocket(s)
		return nil, err
	}

	// This socket is used by a dialer.
	if ursa != nil {
		if !deadline.IsZero() {
			setWriteDeadline(fd, deadline)
		}
		if err = fd.connect(ulsa, ursa); err != nil {
			fd.Close()
			return nil, err
		}
		fd.isConnected = true
		if !deadline.IsZero() {
			setWriteDeadline(fd, noDeadline)
		}
	}

	lsa, _ := syscall.Getsockname(s)
	laddr := toAddr(lsa)
	rsa, _ := syscall.Getpeername(s)
	if rsa == nil {
		rsa = ursa
	}
	raddr := toAddr(rsa)
	fd.setAddr(laddr, raddr)
	return fd, nil
}
<|endoftext|>


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

package sync_test

import (
	"runtime"
	. "sync"
	"sync/atomic"
	"testing"
)

type one int

func (o *one) Increment() {
	*o++
}

func run(once *Once, o *one, c chan bool) {
	once.Do(func() { o.Increment() })
	c <- true
}

func TestOnce(t *testing.T) {
	o := new(one)
	once := new(Once)
	c := make(chan bool)
	const N = 10
	for i := 0; i < N; i++ {
		go run(once, o, c)
	}
	for i := 0; i < N; i++ {
		<-c
	}
	if *o != 1 {
		t.Errorf("once failed: %d is not 1", *o)
	}
}

func BenchmarkOnce(b *testing.B) {
	const CallsPerSched = 1000
	procs := runtime.GOMAXPROCS(-1)
	N := int32(b.N / CallsPerSched)
	var once Once
	f := func() {}
	c := make(chan bool, procs)
	for p := 0; p < procs; p++ {
		go func() {
			for atomic.AddInt32(&N, -1) >= 0 {
				runtime.Gosched()
				for g := 0; g < CallsPerSched; g++ {
					once.Do(f)
				}
			}
			c <- true
		}()
	}
	for p := 0; p < procs; p++ {
		<-c
	}
}
<commit_msg>sync: add Once test with panic Tests behavior of Once when initialization function panics.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sync_test

import (
	"runtime"
	. "sync"
	"sync/atomic"
	"testing"
)

type one int

func (o *one) Increment() {
	*o++
}

func run(t *testing.T, once *Once, o *one, c chan bool) {
	once.Do(func() { o.Increment() })
	if v := *o; v != 1 {
		t.Errorf("once failed inside run: %d is not 1", v)
	}
	c <- true
}

func TestOnce(t *testing.T) {
	o := new(one)
	once := new(Once)
	c := make(chan bool)
	const N = 10
	for i := 0; i < N; i++ {
		go run(t, once, o, c)
	}
	for i := 0; i < N; i++ {
		<-c
	}
	if *o != 1 {
		t.Errorf("once failed outside run: %d is not 1", *o)
	}
}

func TestOncePanic(t *testing.T) {
	once := new(Once)
	for i := 0; i < 2; i++ {
		func() {
			defer func() {
				if recover() == nil {
					t.Fatalf("Once.Do() has not panic'ed")
				}
			}()
			once.Do(func() {
				panic("failed")
			})
		}()
	}
	once.Do(func() {})
	once.Do(func() {
		t.Fatalf("Once called twice")
	})
}

func BenchmarkOnce(b *testing.B) {
	const CallsPerSched = 1000
	procs := runtime.GOMAXPROCS(-1)
	N := int32(b.N / CallsPerSched)
	var once Once
	f := func() {}
	c := make(chan bool, procs)
	for p := 0; p < procs; p++ {
		go func() {
			for atomic.AddInt32(&N, -1) >= 0 {
				runtime.Gosched()
				for g := 0; g < CallsPerSched; g++ {
					once.Do(f)
				}
			}
			c <- true
		}()
	}
	for p := 0; p < procs; p++ {
		<-c
	}
}
<|endoftext|>


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

// Pool is no-op under race detector, so all these tests do not work.
// +build !race

package sync_test

import (
	"runtime"
	"runtime/debug"
	. "sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestPool(t *testing.T) {
	// disable GC so we can control when it happens.
	defer debug.SetGCPercent(debug.SetGCPercent(-1))
	var p Pool
	if p.Get() != nil {
		t.Fatal("expected empty")
	}
	p.Put("a")
	p.Put("b")
	if g := p.Get(); g != "a" {
		t.Fatalf("got %#v; want a", g)
	}
	if g := p.Get(); g != "b" {
		t.Fatalf("got %#v; want b", g)
	}
	if g := p.Get(); g != nil {
		t.Fatalf("got %#v; want nil", g)
	}

	p.Put("c")
	debug.SetGCPercent(100) // to allow following GC to actually run
	runtime.GC()
	if g := p.Get(); g != nil {
		t.Fatalf("got %#v; want nil after GC", g)
	}
}

func TestPoolNew(t *testing.T) {
	// disable GC so we can control when it happens.
	defer debug.SetGCPercent(debug.SetGCPercent(-1))

	i := 0
	p := Pool{
		New: func() interface{} {
			i++
			return i
		},
	}
	if v := p.Get(); v != 1 {
		t.Fatalf("got %v; want 1", v)
	}
	if v := p.Get(); v != 2 {
		t.Fatalf("got %v; want 2", v)
	}
	p.Put(42)
	if v := p.Get(); v != 42 {
		t.Fatalf("got %v; want 42", v)
	}
	if v := p.Get(); v != 3 {
		t.Fatalf("got %v; want 3", v)
	}
}

// Test that Pool does not hold pointers to previously cached
// resources
func TestPoolGC(t *testing.T) {
	var p Pool
	var fin uint32
	const N = 100
	for i := 0; i < N; i++ {
		v := new(string)
		runtime.SetFinalizer(v, func(vv *string) {
			atomic.AddUint32(&fin, 1)
		})
		p.Put(v)
	}
	for i := 0; i < N; i++ {
		p.Get()
	}
	for i := 0; i < 5; i++ {
		runtime.GC()
		time.Sleep(time.Duration(i*100+10) * time.Millisecond)
		// 1 pointer can remain on stack or elsewhere
		if atomic.LoadUint32(&fin) >= N-1 {
			return
		}
	}
	t.Fatalf("only %v out of %v resources are finalized",
		atomic.LoadUint32(&fin), N)
}

func TestPoolStress(t *testing.T) {
	const P = 10
	N := int(1e6)
	if testing.Short() {
		N /= 100
	}
	var p Pool
	done := make(chan bool)
	for i := 0; i < P; i++ {
		go func() {
			var v interface{} = 0
			for j := 0; j < N; j++ {
				if v == nil {
					v = 0
				}
				p.Put(v)
				v = p.Get()
				if v != nil && v.(int) != 0 {
					t.Fatalf("expect 0, got %v", v)
				}
			}
			done <- true
		}()
	}
	for i := 0; i < P; i++ {
		<-done
	}
}

func BenchmarkPool(b *testing.B) {
	var p Pool
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			p.Put(1)
			p.Get()
		}
	})
}

func BenchmarkPoolOverlflow(b *testing.B) {
	var p Pool
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			for b := 0; b < 100; b++ {
				p.Put(1)
			}
			for b := 0; b < 100; b++ {
				p.Get()
			}
		}
	})
}
<commit_msg>sync: fix typo in benchmark name<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Pool is no-op under race detector, so all these tests do not work.
// +build !race

package sync_test

import (
	"runtime"
	"runtime/debug"
	. "sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestPool(t *testing.T) {
	// disable GC so we can control when it happens.
	defer debug.SetGCPercent(debug.SetGCPercent(-1))
	var p Pool
	if p.Get() != nil {
		t.Fatal("expected empty")
	}
	p.Put("a")
	p.Put("b")
	if g := p.Get(); g != "a" {
		t.Fatalf("got %#v; want a", g)
	}
	if g := p.Get(); g != "b" {
		t.Fatalf("got %#v; want b", g)
	}
	if g := p.Get(); g != nil {
		t.Fatalf("got %#v; want nil", g)
	}

	p.Put("c")
	debug.SetGCPercent(100) // to allow following GC to actually run
	runtime.GC()
	if g := p.Get(); g != nil {
		t.Fatalf("got %#v; want nil after GC", g)
	}
}

func TestPoolNew(t *testing.T) {
	// disable GC so we can control when it happens.
	defer debug.SetGCPercent(debug.SetGCPercent(-1))

	i := 0
	p := Pool{
		New: func() interface{} {
			i++
			return i
		},
	}
	if v := p.Get(); v != 1 {
		t.Fatalf("got %v; want 1", v)
	}
	if v := p.Get(); v != 2 {
		t.Fatalf("got %v; want 2", v)
	}
	p.Put(42)
	if v := p.Get(); v != 42 {
		t.Fatalf("got %v; want 42", v)
	}
	if v := p.Get(); v != 3 {
		t.Fatalf("got %v; want 3", v)
	}
}

// Test that Pool does not hold pointers to previously cached
// resources
func TestPoolGC(t *testing.T) {
	var p Pool
	var fin uint32
	const N = 100
	for i := 0; i < N; i++ {
		v := new(string)
		runtime.SetFinalizer(v, func(vv *string) {
			atomic.AddUint32(&fin, 1)
		})
		p.Put(v)
	}
	for i := 0; i < N; i++ {
		p.Get()
	}
	for i := 0; i < 5; i++ {
		runtime.GC()
		time.Sleep(time.Duration(i*100+10) * time.Millisecond)
		// 1 pointer can remain on stack or elsewhere
		if atomic.LoadUint32(&fin) >= N-1 {
			return
		}
	}
	t.Fatalf("only %v out of %v resources are finalized",
		atomic.LoadUint32(&fin), N)
}

func TestPoolStress(t *testing.T) {
	const P = 10
	N := int(1e6)
	if testing.Short() {
		N /= 100
	}
	var p Pool
	done := make(chan bool)
	for i := 0; i < P; i++ {
		go func() {
			var v interface{} = 0
			for j := 0; j < N; j++ {
				if v == nil {
					v = 0
				}
				p.Put(v)
				v = p.Get()
				if v != nil && v.(int) != 0 {
					t.Fatalf("expect 0, got %v", v)
				}
			}
			done <- true
		}()
	}
	for i := 0; i < P; i++ {
		<-done
	}
}

func BenchmarkPool(b *testing.B) {
	var p Pool
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			p.Put(1)
			p.Get()
		}
	})
}

func BenchmarkPoolOverflow(b *testing.B) {
	var p Pool
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			for b := 0; b < 100; b++ {
				p.Put(1)
			}
			for b := 0; b < 100; b++ {
				p.Get()
			}
		}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"go/parser"
	"go/token"
	"os"
	"testing"

	"github.com/Sirupsen/logrus"
	"github.com/stretchr/testify/require"
)

func TestParentPackages(t *testing.T) {
	assert := require.New(t)

	packages := parentPackages("", "github.com/stretchr/testify/require")
	assert.Len(packages, 4)
	assert.True(packages["github.com"])
	assert.True(packages["github.com/stretchr"])
	assert.True(packages["github.com/stretchr/testify"])
	assert.True(packages["github.com/stretchr/testify/require"])
}

func TestParseDir(t *testing.T) {
	assert := require.New(t)

	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, ".", nil, parser.ImportsOnly)
	if err != nil {
		assert.Fail(err.Error())
		return
	}
	for pkgName, pkg := range pkgs {
		fmt.Printf("pkg: %s\n", pkgName)
		if len(pkg.Files) == 0 {
			fmt.Println("  (no files)")
		}
		for name, f := range pkg.Files {
			fmt.Println("  " + name)
			for _, v := range f.Imports {
				fmt.Println("    " + v.Path.Value[1:len(v.Path.Value)-1])
			}
		}
	}
}

func TestListPackages(t *testing.T) {
	assert := require.New(t)
	os.Chdir("../../..")
	for p := range listPackages("github.com/rancher/trash") {
		logrus.Debug(p)
		assert.Equal(3, len(p))
		assert.Contains(p, "github.com/rancher/trash")
		assert.Contains(p, "github.com/rancher/trash/util")
		assert.Contains(p, "github.com/rancher/trash/conf")
	}
}
<commit_msg>fix trash_test.go<commit_after>package main

import (
	"fmt"
	"go/parser"
	"go/token"
	"os"
	"testing"

	"github.com/Sirupsen/logrus"
	"github.com/stretchr/testify/require"
)

func TestParentPackages(t *testing.T) {
	assert := require.New(t)

	packages := parentPackages("", "github.com/stretchr/testify/require")
	assert.Len(packages, 4)
	assert.True(packages["github.com"])
	assert.True(packages["github.com/stretchr"])
	assert.True(packages["github.com/stretchr/testify"])
	assert.True(packages["github.com/stretchr/testify/require"])
}

func TestParseDir(t *testing.T) {
	assert := require.New(t)

	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, ".", nil, parser.ImportsOnly)
	if err != nil {
		assert.Fail(err.Error())
		return
	}
	for pkgName, pkg := range pkgs {
		fmt.Printf("pkg: %s\n", pkgName)
		if len(pkg.Files) == 0 {
			fmt.Println("  (no files)")
		}
		for name, f := range pkg.Files {
			fmt.Println("  " + name)
			for _, v := range f.Imports {
				fmt.Println("    " + v.Path.Value[1:len(v.Path.Value)-1])
			}
		}
	}
}

func TestListPackages(t *testing.T) {
	assert := require.New(t)
	os.Chdir("../../..")
	p := listPackages("github.com/rancher/trash")
	logrus.Debug(p)
	assert.Equal(3, len(p))
	assert.Contains(p, "github.com/rancher/trash")
	assert.Contains(p, "github.com/rancher/trash/util")
	assert.Contains(p, "github.com/rancher/trash/conf")
}
<|endoftext|>


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

package http

import (
	"encoding/json"
	"fmt"
	"reflect"
	"testing"
	"time"
)

var writeSetCookiesTests = []struct {
	Cookie *Cookie
	Raw    string
}{
	{
		&Cookie{Name: "cookie-1", Value: "v$1"},
		"cookie-1=v$1",
	},
	{
		&Cookie{Name: "cookie-2", Value: "two", MaxAge: 3600},
		"cookie-2=two; Max-Age=3600",
	},
	{
		&Cookie{Name: "cookie-3", Value: "three", Domain: ".example.com"},
		"cookie-3=three; Domain=.example.com",
	},
	{
		&Cookie{Name: "cookie-4", Value: "four", Path: "/restricted/"},
		"cookie-4=four; Path=/restricted/",
	},
}

func TestWriteSetCookies(t *testing.T) {
	for i, tt := range writeSetCookiesTests {
		if g, e := tt.Cookie.String(), tt.Raw; g != e {
			t.Errorf("Test %d, expecting:\n%s\nGot:\n%s\n", i, e, g)
			continue
		}
	}
}

type headerOnlyResponseWriter Header

func (ho headerOnlyResponseWriter) Header() Header {
	return Header(ho)
}

func (ho headerOnlyResponseWriter) Write([]byte) (int, error) {
	panic("NOIMPL")
}

func (ho headerOnlyResponseWriter) WriteHeader(int) {
	panic("NOIMPL")
}

func TestSetCookie(t *testing.T) {
	m := make(Header)
	SetCookie(headerOnlyResponseWriter(m), &Cookie{Name: "cookie-1", Value: "one", Path: "/restricted/"})
	SetCookie(headerOnlyResponseWriter(m), &Cookie{Name: "cookie-2", Value: "two", MaxAge: 3600})
	if l := len(m["Set-Cookie"]); l != 2 {
		t.Fatalf("expected %d cookies, got %d", 2, l)
	}
	if g, e := m["Set-Cookie"][0], "cookie-1=one; Path=/restricted/"; g != e {
		t.Errorf("cookie #1: want %q, got %q", e, g)
	}
	if g, e := m["Set-Cookie"][1], "cookie-2=two; Max-Age=3600"; g != e {
		t.Errorf("cookie #2: want %q, got %q", e, g)
	}
}

var addCookieTests = []struct {
	Cookies []*Cookie
	Raw     string
}{
	{
		[]*Cookie{},
		"",
	},
	{
		[]*Cookie{{Name: "cookie-1", Value: "v$1"}},
		"cookie-1=v$1",
	},
	{
		[]*Cookie{
			{Name: "cookie-1", Value: "v$1"},
			{Name: "cookie-2", Value: "v$2"},
			{Name: "cookie-3", Value: "v$3"},
		},
		"cookie-1=v$1; cookie-2=v$2; cookie-3=v$3",
	},
}

func TestAddCookie(t *testing.T) {
	for i, tt := range addCookieTests {
		req, _ := NewRequest("GET", "http://example.com/", nil)
		for _, c := range tt.Cookies {
			req.AddCookie(c)
		}
		if g := req.Header.Get("Cookie"); g != tt.Raw {
			t.Errorf("Test %d:\nwant: %s\n got: %s\n", i, tt.Raw, g)
			continue
		}
	}
}

var readSetCookiesTests = []struct {
	Header  Header
	Cookies []*Cookie
}{
	{
		Header{"Set-Cookie": {"Cookie-1=v$1"}},
		[]*Cookie{{Name: "Cookie-1", Value: "v$1", Raw: "Cookie-1=v$1"}},
	},
	{
		Header{"Set-Cookie": {"NID=99=YsDT5i3E-CXax-; expires=Wed, 23-Nov-2011 01:05:03 GMT; path=/; domain=.google.ch; HttpOnly"}},
		[]*Cookie{{
			Name:       "NID",
			Value:      "99=YsDT5i3E-CXax-",
			Path:       "/",
			Domain:     ".google.ch",
			HttpOnly:   true,
			Expires:    time.Date(2011, 11, 23, 1, 5, 3, 0, time.UTC),
			RawExpires: "Wed, 23-Nov-2011 01:05:03 GMT",
			Raw:        "NID=99=YsDT5i3E-CXax-; expires=Wed, 23-Nov-2011 01:05:03 GMT; path=/; domain=.google.ch; HttpOnly",
		}},
	},
}

func toJSON(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		return fmt.Sprintf("%#v", v)
	}
	return string(b)
}

func TestReadSetCookies(t *testing.T) {
	for i, tt := range readSetCookiesTests {
		for n := 0; n < 2; n++ { // to verify readSetCookies doesn't mutate its input
			c := readSetCookies(tt.Header)
			if !reflect.DeepEqual(c, tt.Cookies) {
				t.Errorf("#%d readSetCookies: have\n%s\nwant\n%s\n", i, toJSON(c), toJSON(tt.Cookies))
				continue
			}
		}
	}
}

var readCookiesTests = []struct {
	Header  Header
	Filter  string
	Cookies []*Cookie
}{
	{
		Header{"Cookie": {"Cookie-1=v$1", "c2=v2"}},
		"",
		[]*Cookie{
			{Name: "Cookie-1", Value: "v$1"},
			{Name: "c2", Value: "v2"},
		},
	},
	{
		Header{"Cookie": {"Cookie-1=v$1", "c2=v2"}},
		"c2",
		[]*Cookie{
			{Name: "c2", Value: "v2"},
		},
	},
	{
		Header{"Cookie": {"Cookie-1=v$1; c2=v2"}},
		"",
		[]*Cookie{
			{Name: "Cookie-1", Value: "v$1"},
			{Name: "c2", Value: "v2"},
		},
	},
	{
		Header{"Cookie": {"Cookie-1=v$1; c2=v2"}},
		"c2",
		[]*Cookie{
			{Name: "c2", Value: "v2"},
		},
	},
}

func TestReadCookies(t *testing.T) {
	for i, tt := range readCookiesTests {
		for n := 0; n < 2; n++ { // to verify readCookies doesn't mutate its input
			c := readCookies(tt.Header, tt.Filter)
			if !reflect.DeepEqual(c, tt.Cookies) {
				t.Errorf("#%d readCookies:\nhave: %s\nwant: %s\n", i, toJSON(c), toJSON(tt.Cookies))
				continue
			}
		}
	}
}
<commit_msg>net/http: some more cookie tests<commit_after>// Copyright 2010 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package http

import (
	"encoding/json"
	"fmt"
	"reflect"
	"testing"
	"time"
)

var writeSetCookiesTests = []struct {
	Cookie *Cookie
	Raw    string
}{
	{
		&Cookie{Name: "cookie-1", Value: "v$1"},
		"cookie-1=v$1",
	},
	{
		&Cookie{Name: "cookie-2", Value: "two", MaxAge: 3600},
		"cookie-2=two; Max-Age=3600",
	},
	{
		&Cookie{Name: "cookie-3", Value: "three", Domain: ".example.com"},
		"cookie-3=three; Domain=.example.com",
	},
	{
		&Cookie{Name: "cookie-4", Value: "four", Path: "/restricted/"},
		"cookie-4=four; Path=/restricted/",
	},
}

func TestWriteSetCookies(t *testing.T) {
	for i, tt := range writeSetCookiesTests {
		if g, e := tt.Cookie.String(), tt.Raw; g != e {
			t.Errorf("Test %d, expecting:\n%s\nGot:\n%s\n", i, e, g)
			continue
		}
	}
}

type headerOnlyResponseWriter Header

func (ho headerOnlyResponseWriter) Header() Header {
	return Header(ho)
}

func (ho headerOnlyResponseWriter) Write([]byte) (int, error) {
	panic("NOIMPL")
}

func (ho headerOnlyResponseWriter) WriteHeader(int) {
	panic("NOIMPL")
}

func TestSetCookie(t *testing.T) {
	m := make(Header)
	SetCookie(headerOnlyResponseWriter(m), &Cookie{Name: "cookie-1", Value: "one", Path: "/restricted/"})
	SetCookie(headerOnlyResponseWriter(m), &Cookie{Name: "cookie-2", Value: "two", MaxAge: 3600})
	if l := len(m["Set-Cookie"]); l != 2 {
		t.Fatalf("expected %d cookies, got %d", 2, l)
	}
	if g, e := m["Set-Cookie"][0], "cookie-1=one; Path=/restricted/"; g != e {
		t.Errorf("cookie #1: want %q, got %q", e, g)
	}
	if g, e := m["Set-Cookie"][1], "cookie-2=two; Max-Age=3600"; g != e {
		t.Errorf("cookie #2: want %q, got %q", e, g)
	}
}

var addCookieTests = []struct {
	Cookies []*Cookie
	Raw     string
}{
	{
		[]*Cookie{},
		"",
	},
	{
		[]*Cookie{{Name: "cookie-1", Value: "v$1"}},
		"cookie-1=v$1",
	},
	{
		[]*Cookie{
			{Name: "cookie-1", Value: "v$1"},
			{Name: "cookie-2", Value: "v$2"},
			{Name: "cookie-3", Value: "v$3"},
		},
		"cookie-1=v$1; cookie-2=v$2; cookie-3=v$3",
	},
}

func TestAddCookie(t *testing.T) {
	for i, tt := range addCookieTests {
		req, _ := NewRequest("GET", "http://example.com/", nil)
		for _, c := range tt.Cookies {
			req.AddCookie(c)
		}
		if g := req.Header.Get("Cookie"); g != tt.Raw {
			t.Errorf("Test %d:\nwant: %s\n got: %s\n", i, tt.Raw, g)
			continue
		}
	}
}

var readSetCookiesTests = []struct {
	Header  Header
	Cookies []*Cookie
}{
	{
		Header{"Set-Cookie": {"Cookie-1=v$1"}},
		[]*Cookie{{Name: "Cookie-1", Value: "v$1", Raw: "Cookie-1=v$1"}},
	},
	{
		Header{"Set-Cookie": {"NID=99=YsDT5i3E-CXax-; expires=Wed, 23-Nov-2011 01:05:03 GMT; path=/; domain=.google.ch; HttpOnly"}},
		[]*Cookie{{
			Name:       "NID",
			Value:      "99=YsDT5i3E-CXax-",
			Path:       "/",
			Domain:     ".google.ch",
			HttpOnly:   true,
			Expires:    time.Date(2011, 11, 23, 1, 5, 3, 0, time.UTC),
			RawExpires: "Wed, 23-Nov-2011 01:05:03 GMT",
			Raw:        "NID=99=YsDT5i3E-CXax-; expires=Wed, 23-Nov-2011 01:05:03 GMT; path=/; domain=.google.ch; HttpOnly",
		}},
	},
	{
		Header{"Set-Cookie": {".ASPXAUTH=7E3AA; expires=Wed, 07-Mar-2012 14:25:06 GMT; path=/; HttpOnly"}},
		[]*Cookie{{
			Name:       ".ASPXAUTH",
			Value:      "7E3AA",
			Path:       "/",
			Expires:    time.Date(2012, 3, 7, 14, 25, 6, 0, time.UTC),
			RawExpires: "Wed, 07-Mar-2012 14:25:06 GMT",
			HttpOnly:   true,
			Raw:        ".ASPXAUTH=7E3AA; expires=Wed, 07-Mar-2012 14:25:06 GMT; path=/; HttpOnly",
		}},
	},
	{
		Header{"Set-Cookie": {"ASP.NET_SessionId=foo; path=/; HttpOnly"}},
		[]*Cookie{{
			Name:     "ASP.NET_SessionId",
			Value:    "foo",
			Path:     "/",
			HttpOnly: true,
			Raw:      "ASP.NET_SessionId=foo; path=/; HttpOnly",
		}},
	},

	// TODO(bradfitz): users have reported seeing this in the
	// wild, but do browsers handle it? RFC 6265 just says "don't
	// do that" (section 3) and then never mentions header folding
	// again.
	// Header{"Set-Cookie": {"ASP.NET_SessionId=foo; path=/; HttpOnly, .ASPXAUTH=7E3AA; expires=Wed, 07-Mar-2012 14:25:06 GMT; path=/; HttpOnly"}},
}

func toJSON(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		return fmt.Sprintf("%#v", v)
	}
	return string(b)
}

func TestReadSetCookies(t *testing.T) {
	for i, tt := range readSetCookiesTests {
		for n := 0; n < 2; n++ { // to verify readSetCookies doesn't mutate its input
			c := readSetCookies(tt.Header)
			if !reflect.DeepEqual(c, tt.Cookies) {
				t.Errorf("#%d readSetCookies: have\n%s\nwant\n%s\n", i, toJSON(c), toJSON(tt.Cookies))
				continue
			}
		}
	}
}

var readCookiesTests = []struct {
	Header  Header
	Filter  string
	Cookies []*Cookie
}{
	{
		Header{"Cookie": {"Cookie-1=v$1", "c2=v2"}},
		"",
		[]*Cookie{
			{Name: "Cookie-1", Value: "v$1"},
			{Name: "c2", Value: "v2"},
		},
	},
	{
		Header{"Cookie": {"Cookie-1=v$1", "c2=v2"}},
		"c2",
		[]*Cookie{
			{Name: "c2", Value: "v2"},
		},
	},
	{
		Header{"Cookie": {"Cookie-1=v$1; c2=v2"}},
		"",
		[]*Cookie{
			{Name: "Cookie-1", Value: "v$1"},
			{Name: "c2", Value: "v2"},
		},
	},
	{
		Header{"Cookie": {"Cookie-1=v$1; c2=v2"}},
		"c2",
		[]*Cookie{
			{Name: "c2", Value: "v2"},
		},
	},
}

func TestReadCookies(t *testing.T) {
	for i, tt := range readCookiesTests {
		for n := 0; n < 2; n++ { // to verify readCookies doesn't mutate its input
			c := readCookies(tt.Header, tt.Filter)
			if !reflect.DeepEqual(c, tt.Cookies) {
				t.Errorf("#%d readCookies:\nhave: %s\nwant: %s\n", i, toJSON(c), toJSON(tt.Cookies))
				continue
			}
		}
	}
}
<|endoftext|>


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

package net

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"os/exec"
	"syscall"
	"testing"
	"time"
)

func TestAcceptIgnoreSomeErrors(t *testing.T) {
	recv := func(ln Listener) (string, error) {
		c, err := ln.Accept()
		if err != nil {
			// Display windows errno in error message.
			operr, ok := err.(*OpError)
			if !ok {
				return "", err
			}
			errno, ok := operr.Err.(syscall.Errno)
			if !ok {
				return "", err
			}
			return "", fmt.Errorf("%v (windows errno=%d)", err, errno)
		}
		defer c.Close()

		b := make([]byte, 100)
		n, err := c.Read(b)
		if err != nil && err != io.EOF {
			return "", err
		}
		return string(b[:n]), nil
	}

	send := func(addr string, data string) error {
		c, err := Dial("tcp", addr)
		if err != nil {
			return err
		}
		defer c.Close()

		b := []byte(data)
		n, err := c.Write(b)
		if err != nil {
			return err
		}
		if n != len(b) {
			return fmt.Errorf(`Only %d chars of string "%s" sent`, n, data)
		}
		return nil
	}

	if envaddr := os.Getenv("GOTEST_DIAL_ADDR"); envaddr != "" {
		// In child process.
		c, err := Dial("tcp", envaddr)
		if err != nil {
			t.Fatalf("Dial failed: %v", err)
		}
		fmt.Printf("sleeping\n")
		time.Sleep(time.Minute) // process will be killed here
		c.Close()
	}

	ln, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("Listen failed: %v", err)
	}
	defer ln.Close()

	// Start child process that connects to our listener.
	cmd := exec.Command(os.Args[0], "-test.run=TestAcceptIgnoreSomeErrors")
	cmd.Env = append(os.Environ(), "GOTEST_DIAL_ADDR="+ln.Addr().String())
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatalf("cmd.StdoutPipe failed: %v", err)
	}
	err = cmd.Start()
	if err != nil {
		t.Fatalf("cmd.Start failed: %v\n%s\n", err)
	}
	outReader := bufio.NewReader(stdout)
	for {
		s, err := outReader.ReadString('\n')
		if err != nil {
			t.Fatalf("reading stdout failed: %v", err)
		}
		if s == "sleeping\n" {
			break
		}
	}
	defer cmd.Wait() // ignore error - we know it is getting killed

	const alittle = 100 * time.Millisecond
	time.Sleep(alittle)
	cmd.Process.Kill() // the only way to trigger the errors
	time.Sleep(alittle)

	// Send second connection data (with delay in a separate goroutine).
	result := make(chan error)
	go func() {
		time.Sleep(alittle)
		err := send(ln.Addr().String(), "abc")
		if err != nil {
			result <- err
		}
		result <- nil
	}()
	defer func() {
		err := <-result
		if err != nil {
			t.Fatalf("send failed: %v", err)
		}
	}()

	// Receive first or second connection.
	s, err := recv(ln)
	if err != nil {
		t.Fatalf("recv failed: %v", err)
	}
	switch s {
	case "":
		// First connection data is received, lets get second connection data.
	case "abc":
		// First connection is lost forever, but that is ok.
		return
	default:
		t.Fatalf(`"%s" received from recv, but "" or "abc" expected`, s)
	}

	// Get second connection data.
	s, err = recv(ln)
	if err != nil {
		t.Fatalf("recv failed: %v", err)
	}
	if s != "abc" {
		t.Fatalf(`"%s" received from recv, but "abc" expected`, s)
	}
}
<commit_msg>net: fix format string in TestAcceptIgnoreSomeErrors<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"os/exec"
	"syscall"
	"testing"
	"time"
)

func TestAcceptIgnoreSomeErrors(t *testing.T) {
	recv := func(ln Listener) (string, error) {
		c, err := ln.Accept()
		if err != nil {
			// Display windows errno in error message.
			operr, ok := err.(*OpError)
			if !ok {
				return "", err
			}
			errno, ok := operr.Err.(syscall.Errno)
			if !ok {
				return "", err
			}
			return "", fmt.Errorf("%v (windows errno=%d)", err, errno)
		}
		defer c.Close()

		b := make([]byte, 100)
		n, err := c.Read(b)
		if err != nil && err != io.EOF {
			return "", err
		}
		return string(b[:n]), nil
	}

	send := func(addr string, data string) error {
		c, err := Dial("tcp", addr)
		if err != nil {
			return err
		}
		defer c.Close()

		b := []byte(data)
		n, err := c.Write(b)
		if err != nil {
			return err
		}
		if n != len(b) {
			return fmt.Errorf(`Only %d chars of string "%s" sent`, n, data)
		}
		return nil
	}

	if envaddr := os.Getenv("GOTEST_DIAL_ADDR"); envaddr != "" {
		// In child process.
		c, err := Dial("tcp", envaddr)
		if err != nil {
			t.Fatalf("Dial failed: %v", err)
		}
		fmt.Printf("sleeping\n")
		time.Sleep(time.Minute) // process will be killed here
		c.Close()
	}

	ln, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("Listen failed: %v", err)
	}
	defer ln.Close()

	// Start child process that connects to our listener.
	cmd := exec.Command(os.Args[0], "-test.run=TestAcceptIgnoreSomeErrors")
	cmd.Env = append(os.Environ(), "GOTEST_DIAL_ADDR="+ln.Addr().String())
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatalf("cmd.StdoutPipe failed: %v", err)
	}
	err = cmd.Start()
	if err != nil {
		t.Fatalf("cmd.Start failed: %v\n", err)
	}
	outReader := bufio.NewReader(stdout)
	for {
		s, err := outReader.ReadString('\n')
		if err != nil {
			t.Fatalf("reading stdout failed: %v", err)
		}
		if s == "sleeping\n" {
			break
		}
	}
	defer cmd.Wait() // ignore error - we know it is getting killed

	const alittle = 100 * time.Millisecond
	time.Sleep(alittle)
	cmd.Process.Kill() // the only way to trigger the errors
	time.Sleep(alittle)

	// Send second connection data (with delay in a separate goroutine).
	result := make(chan error)
	go func() {
		time.Sleep(alittle)
		err := send(ln.Addr().String(), "abc")
		if err != nil {
			result <- err
		}
		result <- nil
	}()
	defer func() {
		err := <-result
		if err != nil {
			t.Fatalf("send failed: %v", err)
		}
	}()

	// Receive first or second connection.
	s, err := recv(ln)
	if err != nil {
		t.Fatalf("recv failed: %v", err)
	}
	switch s {
	case "":
		// First connection data is received, lets get second connection data.
	case "abc":
		// First connection is lost forever, but that is ok.
		return
	default:
		t.Fatalf(`"%s" received from recv, but "" or "abc" expected`, s)
	}

	// Get second connection data.
	s, err = recv(ln)
	if err != nil {
		t.Fatalf("recv failed: %v", err)
	}
	if s != "abc" {
		t.Fatalf(`"%s" received from recv, but "abc" expected`, s)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// +build akaros

package parlib

/*
#include <parlib.h>
#include <uthread.h>
#include <vcore.h>
#include <mcs.h>
#include <futex.h>
*/
import "C"
import (
	"unsafe"
	"errors"
)

var Procinfo *ProcinfoType = (*ProcinfoType)(unsafe.Pointer(uintptr(C.UINFO)))

func Futex(uaddr *int32, op int32, val int32,
           timeout *Timespec, uaddr2 *int32, val3 int32) (ret int32) {
	// For now, akaros futexes don't support uaddr2 or val3, so we
	// just 0 them out.
	uaddr2 = nil;
	val3 = 0;
	// Also, the minimum timout is 1ms, so up it to that if it's too small
	if (timeout != nil) {
		if (timeout.Sec == 0) {
			if (timeout.Nsec < 1000000) {
				timeout.Nsec = 1000000;
			}
		}
	}
	return int32(C.futex((*C.int)(unsafe.Pointer(uaddr)),
	                     C.int(op), C.int(val),
	                     (*C.struct_timespec)(unsafe.Pointer(timeout)),
	                     (*C.int)(unsafe.Pointer(uaddr2)), C.int(val3)))
}

func ProcinfoPackArgs(argv []*byte, envp []*byte) (pi ProcinfoType, err error) {
	p_pi := (*_Ctype_procinfo_t)(unsafe.Pointer(&pi))
    p_argv := (**_Ctype_char)(unsafe.Pointer(&argv[0]))
    p_envp := (**_Ctype_char)(unsafe.Pointer(&envp[0]))

	__err := C.procinfo_pack_args(p_pi, p_argv, p_envp)
	if __err == -1 {
		err = nil
	} else {
		err = errors.New("ProcinfoPackArgs: error packing argv and envp")
	}
	return pi, err
}

<commit_msg>Explicitly set the cgoflags needed by the parlib package<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// +build akaros

package parlib

/*
#include <parlib.h>
#include <uthread.h>
#include <vcore.h>
#include <mcs.h>
#include <futex.h>

#cgo akaros LDFLAGS: -lpthread -lbenchutil -lm
*/
import "C"
import (
	"unsafe"
	"errors"
)

var Procinfo *ProcinfoType = (*ProcinfoType)(unsafe.Pointer(uintptr(C.UINFO)))

func Futex(uaddr *int32, op int32, val int32,
           timeout *Timespec, uaddr2 *int32, val3 int32) (ret int32) {
	// For now, akaros futexes don't support uaddr2 or val3, so we
	// just 0 them out.
	uaddr2 = nil;
	val3 = 0;
	// Also, the minimum timout is 1ms, so up it to that if it's too small
	if (timeout != nil) {
		if (timeout.Sec == 0) {
			if (timeout.Nsec < 1000000) {
				timeout.Nsec = 1000000;
			}
		}
	}
	return int32(C.futex((*C.int)(unsafe.Pointer(uaddr)),
	                     C.int(op), C.int(val),
	                     (*C.struct_timespec)(unsafe.Pointer(timeout)),
	                     (*C.int)(unsafe.Pointer(uaddr2)), C.int(val3)))
}

func ProcinfoPackArgs(argv []*byte, envp []*byte) (pi ProcinfoType, err error) {
	p_pi := (*_Ctype_procinfo_t)(unsafe.Pointer(&pi))
    p_argv := (**_Ctype_char)(unsafe.Pointer(&argv[0]))
    p_envp := (**_Ctype_char)(unsafe.Pointer(&envp[0]))

	__err := C.procinfo_pack_args(p_pi, p_argv, p_envp)
	if __err == -1 {
		err = nil
	} else {
		err = errors.New("ProcinfoPackArgs: error packing argv and envp")
	}
	return pi, err
}

<|endoftext|>


<commit_before>//// Copyright 2015 Spring Rain Software Compnay LTD. All Rights Reserved.
//// Licensed under the MIT (MIT-LICENSE.txt) license.
package proxy

import (
	"fmt"
	"github.com/wfxiang08/cyutils/utils/atomic2"
	log "github.com/wfxiang08/cyutils/utils/rolling_log"
	topozk "github.com/wfxiang08/go-zookeeper/zk"
	"github.com/wfxiang08/go_thrift/thrift"
	"github.com/wfxiang08/thrift_rpc_base/rpc_utils"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

type Server interface {
	Run()
}
type ServerFactorory func(config *ServiceConfig) Server
type ProxyFactorory func(config *ProxyConfig) Server

//
// Thrift Server的参数
//
type ThriftRpcServer struct {
	ZkAddr          string
	ProductName     string
	ServiceName     string
	FrontendAddr    string
	Topo            *Topology
	Processor       thrift.TProcessor
	Verbose         bool
	lastRequestTime atomic2.Int64
	config          *ServiceConfig
}

func NewThriftRpcServer(config *ServiceConfig, processor thrift.TProcessor) *ThriftRpcServer {
	log.Printf("FrontAddr: %s\n", Magenta(config.FrontendAddr))

	return &ThriftRpcServer{
		config:       config,
		ZkAddr:       config.ZkAddr,
		ProductName:  config.ProductName,
		ServiceName:  config.Service,
		FrontendAddr: config.FrontendAddr,
		Verbose:      config.Verbose,
		Processor:    processor,
	}

}

//
// 根据前端的地址生成服务的id
// 例如: 127.0.0.1:5555 --> 127_0_0_1_5555
// "/Users/feiwang/gowork/src/github.com/wfxiang08/rpc_proxyiplocation/aa.sock"
func GetServiceIdentity(frontendAddr string) string {

	fid := strings.Replace(frontendAddr, ".", "_", -1)
	fid = strings.Replace(fid, ":", "_", -1)
	fid = strings.Replace(fid, "/", "", -1)
	if len(fid) > 20 {
		// unix domain socket
		fid = fid[len(fid)-20 : len(fid)]
	}
	return fid

}

//
// 去ZK注册当前的Service
//
func RegisterService(serviceName, frontendAddr, serviceId string, topo *Topology, evtExit chan interface{},
	workDir string, codeUrlVerion string, state *atomic2.Bool, stateChan chan bool) *ServiceEndpoint {

	// 1. 准备数据
	// 记录Service Endpoint的信息
	servicePath := topo.ProductServicePath(serviceName)

	// 确保东西存在
	if ok, _ := topo.Exist(servicePath); !ok {
		topo.CreateDir(servicePath)
	}

	// 用来从zookeeper获取事件
	evtbus := make(chan interface{})

	// 2. 将信息添加到Zk中, 并且监控Zk的状态(如果添加失败会怎么样?)
	endpoint := NewServiceEndpoint(serviceName, serviceId, frontendAddr, workDir, codeUrlVerion)

	// deployPath

	// 为了保证Add是干净的，需要先删除，保证自己才是Owner
	endpoint.DeleteServiceEndpoint(topo)

	// 如果没有状态，或状态为true, 则上线
	if state == nil || state.Get() {
		endpoint.AddServiceEndpoint(topo)
	}

	go func() {

		for true {
			// Watch的目的是监控: 当前的zk session的状态, 如果session出现异常，则重新注册
			_, err := topo.WatchNode(servicePath, evtbus)

			if err == nil {
				// 如果成功添加Watch, 则等待退出，或者ZK Event
				select {
				case <-evtExit:
					return
				case <-stateChan:
					// 如何状态变化(则重新注册)
					endpoint.DeleteServiceEndpoint(topo)
					if state == nil || state.Get() {
						endpoint.AddServiceEndpoint(topo)
					}
				case e := <-evtbus:
					event := e.(topozk.Event)
					if event.State == topozk.StateExpired ||
						event.Type == topozk.EventNotWatching {
						// Session过期了，则需要删除之前的数据，
						// 因为当前的session不是之前的数据的Owner
						endpoint.DeleteServiceEndpoint(topo)

						if state == nil || state.Get() {
							endpoint.AddServiceEndpoint(topo)
						}

					}
				}

			} else {
				// 如果添加Watch失败，则等待退出，或等待timer接触，进行下一次尝试
				timer := time.NewTimer(time.Second * time.Duration(10))
				log.ErrorErrorf(err, "Reg Failed, Wait 10 seconds...., %v", err)
				select {
				case <-evtExit:
					return
				case <-timer.C:
					// pass
				}
			}

		}
	}()
	return endpoint
}

//
// 后端如何处理一个Request
//
func (p *ThriftRpcServer) Dispatch(r *Request) error {
	transport := NewTMemoryBufferWithBuf(r.Request.Data)
	ip := thrift.NewTBinaryProtocolTransport(transport)

	slice := getSlice(0, DEFAULT_SLICE_LEN)
	transport = NewTMemoryBufferWithBuf(slice)
	op := thrift.NewTBinaryProtocolTransport(transport)
	p.Processor.Process(ip, op)

	r.Response.Data = transport.Bytes()

	_, _, seqId, _ := DecodeThriftTypIdSeqId(r.Response.Data)

	log.Debugf("SeqId: %d vs. %d, Dispatch Over", r.Request.SeqId, seqId)
	//	// 如果transport重新分配了内存，则立即归还slice
	//	if cap(r.Response.Data) != DEFAULT_SLICE_LEN {
	//		returnSlice(slice)
	//	}
	return nil
}

func (p *ThriftRpcServer) Run() {
	//	// 1. 创建到zk的连接
	// 有条件地做服务注册
	registerService := !p.config.StandAlone // 不独立运行则注册服务
	if registerService {
		p.Topo = NewTopology(p.ProductName, p.ZkAddr)
	}

	// 127.0.0.1:5555 --> 127_0_0_1:5555
	lbServiceName := GetServiceIdentity(p.FrontendAddr)

	exitSignal := make(chan os.Signal, 1)
	signal.Notify(exitSignal, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)

	// syscall.SIGKILL
	// kill -9 pid
	// kill -s SIGKILL pid 还是留给运维吧

	if len(p.config.FalconClient) > 0 {
		StartTicker(p.config.FalconClient, p.ServiceName)
	}

	// 初始状态为不上线
	var state atomic2.Bool
	state.Set(false)
	stateChan := make(chan bool)

	// 注册服务
	evtExit := make(chan interface{})

	var endpoint *ServiceEndpoint = nil
	if registerService {
		endpoint = RegisterService(p.ServiceName, p.FrontendAddr, lbServiceName,
			p.Topo, evtExit, p.config.WorkDir, p.config.CodeUrlVersion,
			&state, stateChan)
	}

	// 3. 读取"前端"的配置
	var transport thrift.TServerTransport
	var err error

	isUnixDomain := false
	// 127.0.0.1:9999(以:区分不同的类型)
	if !strings.Contains(p.FrontendAddr, ":") {
		if rpc_utils.FileExist(p.FrontendAddr) {
			os.Remove(p.FrontendAddr)
		}
		transport, err = rpc_utils.NewTServerUnixDomain(p.FrontendAddr)
		isUnixDomain = true
	} else {
		transport, err = thrift.NewTServerSocket(p.FrontendAddr)
	}

	if err != nil {
		log.ErrorErrorf(err, Red("Server Socket Create Failed: %v"), err)
		panic(fmt.Sprintf("Invalid FrontendAddr: %s", p.FrontendAddr))
	}

	err = transport.Listen()
	if err != nil {
		log.ErrorErrorf(err, Red("Server Socket Open Failed: %v"), err)
		panic(fmt.Sprintf("Server Socket Open Failed: %s", p.FrontendAddr))
	}

	ch := make(chan thrift.TTransport, 4096)
	defer close(ch)

	// 强制退出? TODO: Graceful退出
	go func() {
		<-exitSignal
		evtExit <- true
		log.Info(Magenta("Receive Exit Signals...."))
		if registerService {
			endpoint.DeleteServiceEndpoint(p.Topo)
		} else {
			// 如果没有注册到后端服务器，则直接不再接受新请求
			transport.Interrupt()
		}

		// 等待
		start := time.Now().Unix()
		for true {
			// 如果5s内没有接受到新的请求了，则退出
			now := time.Now().Unix()
			if now-p.lastRequestTime.Get() > 5 {
				log.Info(Red("Graceful Exit..."))
				break
			} else {
				log.Printf(Cyan("Sleeping %d seconds\n"), now-start)
				time.Sleep(time.Second)
			}
		}
		if registerService {
			transport.Interrupt()
		}
		transport.Close()
	}()

	go func() {
		var address string
		for c := range ch {
			// 为每个Connection建立一个Session
			socket, ok := c.(rpc_utils.SocketAddr)

			if ok {
				if isUnixDomain {
					address = p.FrontendAddr
				} else {
					address = socket.Addr().String()
				}
			} else {
				address = "unknow"
			}

			// Session独立处理自己的请求
			if registerService {
				x := NewNonBlockSession(c, address, p.Verbose, &p.lastRequestTime)
				go x.Serve(p, 1000)
			} else {
				go func() {
					// 直接处理请求，一次性的
					ip := thrift.NewTBinaryProtocolTransport(thrift.NewTFramedTransport(c))
					op := thrift.NewTBinaryProtocolTransport(thrift.NewTFramedTransport(c))
					p.Processor.Process(ip, op)
					c.Close()
				}()
			}
		}
	}()

	// 准备上线服务
	state.Set(true)
	if registerService {
		stateChan <- true
	}

	log.Printf("Begin to accept requests...")
	// Accept什么时候出错，出错之后如何处理呢?
	for {
		c, err := transport.Accept()
		if err != nil {
			break
		} else {
			ch <- c
		}
	}
}
<commit_msg>rpc service直接对外服务<commit_after>//// Copyright 2015 Spring Rain Software Compnay LTD. All Rights Reserved.
//// Licensed under the MIT (MIT-LICENSE.txt) license.
package proxy

import (
	"fmt"
	"github.com/wfxiang08/cyutils/utils/atomic2"
	log "github.com/wfxiang08/cyutils/utils/rolling_log"
	topozk "github.com/wfxiang08/go-zookeeper/zk"
	"github.com/wfxiang08/go_thrift/thrift"
	"github.com/wfxiang08/thrift_rpc_base/rpc_utils"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

type Server interface {
	Run()
}
type ServerFactorory func(config *ServiceConfig) Server
type ProxyFactorory func(config *ProxyConfig) Server

//
// Thrift Server的参数
//
type ThriftRpcServer struct {
	ZkAddr          string
	ProductName     string
	ServiceName     string
	FrontendAddr    string
	Topo            *Topology
	Processor       thrift.TProcessor
	Verbose         bool
	lastRequestTime atomic2.Int64
	config          *ServiceConfig
}

func NewThriftRpcServer(config *ServiceConfig, processor thrift.TProcessor) *ThriftRpcServer {
	log.Printf("FrontAddr: %s\n", Magenta(config.FrontendAddr))

	return &ThriftRpcServer{
		config:       config,
		ZkAddr:       config.ZkAddr,
		ProductName:  config.ProductName,
		ServiceName:  config.Service,
		FrontendAddr: config.FrontendAddr,
		Verbose:      config.Verbose,
		Processor:    processor,
	}

}

//
// 根据前端的地址生成服务的id
// 例如: 127.0.0.1:5555 --> 127_0_0_1_5555
// "/Users/feiwang/gowork/src/github.com/wfxiang08/rpc_proxyiplocation/aa.sock"
func GetServiceIdentity(frontendAddr string) string {

	fid := strings.Replace(frontendAddr, ".", "_", -1)
	fid = strings.Replace(fid, ":", "_", -1)
	fid = strings.Replace(fid, "/", "", -1)
	if len(fid) > 20 {
		// unix domain socket
		fid = fid[len(fid)-20 : len(fid)]
	}
	return fid

}

//
// 去ZK注册当前的Service
//
func RegisterService(serviceName, frontendAddr, serviceId string, topo *Topology, evtExit chan interface{},
	workDir string, codeUrlVerion string, state *atomic2.Bool, stateChan chan bool) *ServiceEndpoint {

	// 1. 准备数据
	// 记录Service Endpoint的信息
	servicePath := topo.ProductServicePath(serviceName)

	// 确保东西存在
	if ok, _ := topo.Exist(servicePath); !ok {
		topo.CreateDir(servicePath)
	}

	// 用来从zookeeper获取事件
	evtbus := make(chan interface{})

	// 2. 将信息添加到Zk中, 并且监控Zk的状态(如果添加失败会怎么样?)
	endpoint := NewServiceEndpoint(serviceName, serviceId, frontendAddr, workDir, codeUrlVerion)

	// deployPath

	// 为了保证Add是干净的，需要先删除，保证自己才是Owner
	endpoint.DeleteServiceEndpoint(topo)

	// 如果没有状态，或状态为true, 则上线
	if state == nil || state.Get() {
		endpoint.AddServiceEndpoint(topo)
	}

	go func() {

		for true {
			// Watch的目的是监控: 当前的zk session的状态, 如果session出现异常，则重新注册
			_, err := topo.WatchNode(servicePath, evtbus)

			if err == nil {
				// 如果成功添加Watch, 则等待退出，或者ZK Event
				select {
				case <-evtExit:
					return
				case <-stateChan:
					// 如何状态变化(则重新注册)
					endpoint.DeleteServiceEndpoint(topo)
					if state == nil || state.Get() {
						endpoint.AddServiceEndpoint(topo)
					}
				case e := <-evtbus:
					event := e.(topozk.Event)
					if event.State == topozk.StateExpired ||
						event.Type == topozk.EventNotWatching {
						// Session过期了，则需要删除之前的数据，
						// 因为当前的session不是之前的数据的Owner
						endpoint.DeleteServiceEndpoint(topo)

						if state == nil || state.Get() {
							endpoint.AddServiceEndpoint(topo)
						}

					}
				}

			} else {
				// 如果添加Watch失败，则等待退出，或等待timer接触，进行下一次尝试
				timer := time.NewTimer(time.Second * time.Duration(10))
				log.ErrorErrorf(err, "Reg Failed, Wait 10 seconds...., %v", err)
				select {
				case <-evtExit:
					return
				case <-timer.C:
					// pass
				}
			}

		}
	}()
	return endpoint
}

//
// 后端如何处理一个Request
//
func (p *ThriftRpcServer) Dispatch(r *Request) error {
	transport := NewTMemoryBufferWithBuf(r.Request.Data)
	ip := thrift.NewTBinaryProtocolTransport(transport)

	slice := getSlice(0, DEFAULT_SLICE_LEN)
	transport = NewTMemoryBufferWithBuf(slice)
	op := thrift.NewTBinaryProtocolTransport(transport)
	p.Processor.Process(ip, op)

	r.Response.Data = transport.Bytes()

	_, _, seqId, _ := DecodeThriftTypIdSeqId(r.Response.Data)

	log.Debugf("SeqId: %d vs. %d, Dispatch Over", r.Request.SeqId, seqId)
	//	// 如果transport重新分配了内存，则立即归还slice
	//	if cap(r.Response.Data) != DEFAULT_SLICE_LEN {
	//		returnSlice(slice)
	//	}
	return nil
}

func (p *ThriftRpcServer) Run() {
	//	// 1. 创建到zk的连接
	// 有条件地做服务注册
	registerService := !p.config.StandAlone // 不独立运行则注册服务
	if registerService {
		p.Topo = NewTopology(p.ProductName, p.ZkAddr)
	}

	// 127.0.0.1:5555 --> 127_0_0_1:5555
	lbServiceName := GetServiceIdentity(p.FrontendAddr)

	exitSignal := make(chan os.Signal, 1)
	signal.Notify(exitSignal, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)

	// syscall.SIGKILL
	// kill -9 pid
	// kill -s SIGKILL pid 还是留给运维吧

	if len(p.config.FalconClient) > 0 {
		StartTicker(p.config.FalconClient, p.ServiceName)
	}

	// 初始状态为不上线
	var state atomic2.Bool
	state.Set(false)
	stateChan := make(chan bool)

	// 注册服务
	evtExit := make(chan interface{})

	var endpoint *ServiceEndpoint = nil
	if registerService {
		endpoint = RegisterService(p.ServiceName, p.FrontendAddr, lbServiceName,
			p.Topo, evtExit, p.config.WorkDir, p.config.CodeUrlVersion,
			&state, stateChan)
	}

	// 3. 读取"前端"的配置
	var transport thrift.TServerTransport
	var err error

	isUnixDomain := false
	// 127.0.0.1:9999(以:区分不同的类型)
	if !strings.Contains(p.FrontendAddr, ":") {
		if rpc_utils.FileExist(p.FrontendAddr) {
			os.Remove(p.FrontendAddr)
		}
		transport, err = rpc_utils.NewTServerUnixDomain(p.FrontendAddr)
		isUnixDomain = true
	} else {
		transport, err = thrift.NewTServerSocket(p.FrontendAddr)
	}

	if err != nil {
		log.ErrorErrorf(err, Red("Server Socket Create Failed: %v"), err)
		panic(fmt.Sprintf("Invalid FrontendAddr: %s", p.FrontendAddr))
	}

	err = transport.Listen()
	if err != nil {
		log.ErrorErrorf(err, Red("Server Socket Open Failed: %v"), err)
		panic(fmt.Sprintf("Server Socket Open Failed: %s", p.FrontendAddr))
	}

	ch := make(chan thrift.TTransport, 4096)
	defer close(ch)

	// 强制退出? TODO: Graceful退出
	go func() {
		<-exitSignal
		log.Info(Magenta("Receive Exit Signals...."))

		if registerService {
			evtExit <- true
		}
		if registerService {
			endpoint.DeleteServiceEndpoint(p.Topo)
		} else {
			// 如果没有注册到后端服务器，则直接不再接受新请求
			transport.Interrupt()
		}

		// 等待
		start := time.Now().Unix()
		for true {
			// 如果5s内没有接受到新的请求了，则退出
			now := time.Now().Unix()
			if now-p.lastRequestTime.Get() > 5 {
				log.Info(Red("Graceful Exit..."))
				break
			} else {
				log.Printf(Cyan("Sleeping %d seconds\n"), now-start)
				time.Sleep(time.Second)
			}
		}

		// 确定处理完毕数据之后，关闭transport
		if registerService {
			transport.Interrupt()
		}
		transport.Close()
	}()

	go func() {
		var address string
		for c := range ch {
			// 为每个Connection建立一个Session
			socket, ok := c.(rpc_utils.SocketAddr)

			if ok {
				if isUnixDomain {
					address = p.FrontendAddr
				} else {
					address = socket.Addr().String()
				}
			} else {
				address = "unknow"
			}

			// Session独立处理自己的请求
			if registerService {
				x := NewNonBlockSession(c, address, p.Verbose, &p.lastRequestTime)
				go x.Serve(p, 1000)
			} else {
				go func() {
					// 直接处理请求，一次性的
					ip := thrift.NewTBinaryProtocolTransport(thrift.NewTFramedTransport(c))
					op := thrift.NewTBinaryProtocolTransport(thrift.NewTFramedTransport(c))
					p.Processor.Process(ip, op)
					c.Close()
				}()
			}
		}
	}()

	// 准备上线服务
	state.Set(true)
	if registerService {
		stateChan <- true
	}

	log.Printf("Begin to accept requests...")
	// Accept什么时候出错，出错之后如何处理呢?
	for {
		c, err := transport.Accept()
		if err != nil {
			break
		} else {
			ch <- c
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tos

import (
	"io"
	"io/ioutil"
	"os"
)

// ChownByName is a wrapper around ChownId that allows changing user and group by name.
func ChownByName(filePath, usr, grp string) error {
	var uid, gid int
	var err error

	if uid, err = GetUid(usr); err != nil {
		return err
	}

	if gid, err = GetGid(grp); err != nil {
		return err
	}

	return Chown(filePath, uid, gid)
}

// ChownId is a wrapper around os.Chown that allows changing user and group
// recursively if given a directory.
func Chown(filePath string, uid, gid int) error {
	stat, err := os.Lstat(filePath)
	if err != nil {
		return err
	}

	if stat.IsDir() {
		files, err := ioutil.ReadDir(filePath)
		if err != nil {
			return err
		}
		for _, file := range files {
			if err := Chown(filePath+"/"+file.Name(), uid, gid); err != nil {
				return err
			}
		}
	}

	if stat.Mode()&os.ModeSymlink != 0 {
		// TODO: os.Chown fails on symlinks
		return nil
	}

	return os.Chown(filePath, uid, gid)
}

// Chmod is a wrapper around os.Chmod that allows changing rights recursively
// if a directory is given.
func Chmod(filePath string, mode os.FileMode) error {
	stat, err := os.Lstat(filePath)
	if err != nil {
		return err
	}

	if stat.IsDir() {
		files, err := ioutil.ReadDir(filePath)
		if err != nil {
			return err
		}
		for _, file := range files {
			if err := Chmod(filePath+"/"+file.Name(), mode); err != nil {
				return err
			}
		}

		// Set executable rights for folders if read or write is allowed
		execRights := 0
		if mode&0600 != 0 {
			execRights |= 0100
		}
		if mode&0060 != 0 {
			execRights |= 0010
		}
		if mode&0006 != 0 {
			execRights |= 0001
		}

		return os.Chmod(filePath, mode|os.FileMode(execRights))
	}

	if stat.Mode()&os.ModeSymlink != 0 {
		// TODO: os.Chmod fails on symlinks
		return nil
	}

	return os.Chmod(filePath, mode)
}

// Copy is a file copy helper. Files will be copied to their destination,
// overwriting existing files. Already existing files that are not part of the
// copy process will not be touched. If source is a directory it is walked
// recursively. Non-existing folders in dest will be created.
// Copy returns upon the first error encountered. In-between results will not
// be rolled back.
func Copy(dest, source string) error {
	srcStat, err := os.Lstat(source)
	if err != nil {
		return err
	}

	if srcStat.IsDir() {
		files, err := ioutil.ReadDir(source)
		if err != nil {
			return err
		}

		if err := os.MkdirAll(dest, srcStat.Mode()); err != nil && !os.IsExist(err) {
			return err
		}

		for _, file := range files {
			if err := Copy(dest+"/"+file.Name(), source+"/"+file.Name()); err != nil {
				return err
			}
		}
		return nil // ### return, copy done ###
	}

	// Symlink handling
	if srcStat.Mode()&os.ModeSymlink != 0 {
		target, err := os.Readlink(source)
		if err != nil {
			return err
		}
		return os.Symlink(target, dest) // ### return, copy done ###
	}

	srcFile, err := os.Open(source)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	dstFile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer dstFile.Close()

	if _, err = io.Copy(dstFile, srcFile); err != nil {
		return err
	}

	return os.Chmod(dest, srcStat.Mode())
}
<commit_msg>Added tos.IsSymlink<commit_after>// Copyright 2015-2016 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tos

import (
	"io"
	"io/ioutil"
	"os"
)

// ChownByName is a wrapper around ChownId that allows changing user and group by name.
func ChownByName(filePath, usr, grp string) error {
	var uid, gid int
	var err error

	if uid, err = GetUid(usr); err != nil {
		return err
	}

	if gid, err = GetGid(grp); err != nil {
		return err
	}

	return Chown(filePath, uid, gid)
}

// ChownId is a wrapper around os.Chown that allows changing user and group
// recursively if given a directory.
func Chown(filePath string, uid, gid int) error {
	stat, err := os.Lstat(filePath)
	if err != nil {
		return err
	}

	if stat.IsDir() {
		files, err := ioutil.ReadDir(filePath)
		if err != nil {
			return err
		}
		for _, file := range files {
			if err := Chown(filePath+"/"+file.Name(), uid, gid); err != nil {
				return err
			}
		}
	}

	if stat.Mode()&os.ModeSymlink != 0 {
		// TODO: os.Chown fails on symlinks
		return nil
	}

	return os.Chown(filePath, uid, gid)
}

// Chmod is a wrapper around os.Chmod that allows changing rights recursively
// if a directory is given.
func Chmod(filePath string, mode os.FileMode) error {
	stat, err := os.Lstat(filePath)
	if err != nil {
		return err
	}

	if stat.IsDir() {
		files, err := ioutil.ReadDir(filePath)
		if err != nil {
			return err
		}
		for _, file := range files {
			if err := Chmod(filePath+"/"+file.Name(), mode); err != nil {
				return err
			}
		}

		// Set executable rights for folders if read or write is allowed
		execRights := 0
		if mode&0600 != 0 {
			execRights |= 0100
		}
		if mode&0060 != 0 {
			execRights |= 0010
		}
		if mode&0006 != 0 {
			execRights |= 0001
		}

		return os.Chmod(filePath, mode|os.FileMode(execRights))
	}

	if stat.Mode()&os.ModeSymlink != 0 {
		// TODO: os.Chmod fails on symlinks
		return nil
	}

	return os.Chmod(filePath, mode)
}

// IsSymlink returns true if a file is a symlink
func IsSymlink(file string) (bool, error) {
	fileStat, err := os.Lstat(file)
	if err != nil {
		return false, err
	}

	return fileStat.Mode()&os.ModeSymlink != 0, nil
}

// Copy is a file copy helper. Files will be copied to their destination,
// overwriting existing files. Already existing files that are not part of the
// copy process will not be touched. If source is a directory it is walked
// recursively. Non-existing folders in dest will be created.
// Copy returns upon the first error encountered. In-between results will not
// be rolled back.
func Copy(dest, source string) error {
	srcStat, err := os.Lstat(source)
	if err != nil {
		return err
	}

	if srcStat.IsDir() {
		files, err := ioutil.ReadDir(source)
		if err != nil {
			return err
		}

		if err := os.MkdirAll(dest, srcStat.Mode()); err != nil && !os.IsExist(err) {
			return err
		}

		for _, file := range files {
			if err := Copy(dest+"/"+file.Name(), source+"/"+file.Name()); err != nil {
				return err
			}
		}
		return nil // ### return, copy done ###
	}

	// Symlink handling
	if srcStat.Mode()&os.ModeSymlink != 0 {
		target, err := os.Readlink(source)
		if err != nil {
			return err
		}
		return os.Symlink(target, dest) // ### return, copy done ###
	}

	srcFile, err := os.Open(source)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	dstFile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer dstFile.Close()

	if _, err = io.Copy(dstFile, srcFile); err != nil {
		return err
	}

	return os.Chmod(dest, srcStat.Mode())
}
<|endoftext|>


<commit_before>package log

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"time"

	"github.com/fatih/camelcase"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

// Logger is a helper for emitting our grpc API logs
type Logger interface {
	Log(request interface{}, response interface{}, err error, duration time.Duration)
	LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int)
}

type logger struct {
	*logrus.Entry
}

// NewLogger creates a new logger
func NewLogger(service string) Logger {
	l := logrus.New()
	l.Formatter = new(prettyFormatter)
	return &logger{
		l.WithFields(logrus.Fields{"service": service}),
	}
}

// Helper function used to log requests and responses from our GRPC method
// implementations
func (l *logger) Log(request interface{}, response interface{}, err error, duration time.Duration) {
	state := "started"
	if err != nil {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.ErrorLevel, 4)
		state = "errored"
	} else {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.InfoLevel, 4)
	}
	if duration.Seconds() > 0 {
		state = "finished"
	}
	l.ReportMetric(state, duration)
}

func (l *logger) ReportMetric(state string, duration time.Duration) {
	depth := 4
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]
	entry := l.WithFields(logrus.Fields{"method": method})
	fmt.Printf("GOING TO REPORT STATS FOR METHOD: %v\n", method)
	fmt.Printf("state: %v, duration: %v\n", state, duration)

	var tokens []string
	for _, token := range camelcase.Split(method) {
		tokens = append(tokens, strings.ToLower(token))
	}
	rootStatName := strings.Join(tokens, "_")

	bucketFactor := 2.0
	bucketCount := 20 // Which makes the max bucket 2^20 seconds or ~12 days in size
	runTime := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: "pachyderm",
			Subsystem: "pachd",
			Name:      fmt.Sprintf("%v_time", rootStatName),
			Help:      fmt.Sprintf("Run time of %v", method),
			Buckets:   prometheus.ExponentialBuckets(1.0, bucketFactor, bucketCount),
		},
		[]string{
			"state", // Since both finished and errored datums can have proc times
		},
	)
	if err := prometheus.Register(runTime); err != nil {
		l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
	}
	if hist, err := runTime.GetMetricWithLabelValues(state); err != nil {
		l.LogAtLevel(entry, logrus.WarnLevel, "failed to get histogram w labels: state (%v) with error %v", state, err)
	} else {
		hist.Observe(duration.Seconds())
	}

	secondsCount := prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace: "pachyderm",
			Subsystem: "pachd",
			Name:      fmt.Sprintf("%v_seconds_count", rootStatName),
			Help:      fmt.Sprintf("cumulative number of seconds spent in %v", method),
		},
	)
	if err := prometheus.Register(secondsCount); err != nil {
		l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
	}
	secondsCount.Add(duration.Seconds())

}

func (l *logger) LogAtLevel(entry *logrus.Entry, level logrus.Level, args ...interface{}) {
	switch level {
	case logrus.PanicLevel:
		entry.Panic(args)
	case logrus.FatalLevel:
		entry.Fatal(args)
	case logrus.ErrorLevel:
		entry.Error(args)
	case logrus.WarnLevel:
		entry.Warn(args)
	case logrus.InfoLevel:
		entry.Info(args)
	case logrus.DebugLevel:
		entry.Debug(args)
	}
}

func (l *logger) LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int) {
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]

	fields := logrus.Fields{
		"method":  method,
		"request": request,
	}
	if response != nil {
		fields["response"] = response
	}
	if err != nil {
		// "err" itself might be a code or even an empty struct
		fields["error"] = err.Error()
	}
	if duration > 0 {
		fields["duration"] = duration
	}
	l.LogAtLevel(l.WithFields(fields), level)
}

type prettyFormatter struct{}

func (f *prettyFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	serialized := []byte(
		fmt.Sprintf(
			"%v %v ",
			entry.Time.Format(logrus.DefaultTimestampFormat),
			strings.ToUpper(entry.Level.String()),
		),
	)
	if entry.Data["service"] != nil {
		serialized = append(serialized, []byte(fmt.Sprintf("%v.%v ", entry.Data["service"], entry.Data["method"]))...)
	}
	if len(entry.Data) > 2 {
		delete(entry.Data, "service")
		delete(entry.Data, "method")
		if entry.Data["duration"] != nil {
			entry.Data["duration"] = entry.Data["duration"].(time.Duration).Seconds()
		}
		data, err := json.Marshal(entry.Data)
		if err != nil {
			return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
		}
		serialized = append(serialized, []byte(string(data))...)
		serialized = append(serialized, ' ')
	}

	serialized = append(serialized, []byte(entry.Message)...)
	serialized = append(serialized, '\n')
	return serialized, nil
}
<commit_msg>Keep track of prometheus metrics that have already been registered<commit_after>package log

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"time"

	"github.com/fatih/camelcase"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

// Logger is a helper for emitting our grpc API logs
type Logger interface {
	Log(request interface{}, response interface{}, err error, duration time.Duration)
	LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int)
}

type logger struct {
	*logrus.Entry
	histogram map[string]*prometheus.HistogramVec
	counter   map[string]prometheus.Counter
}

// NewLogger creates a new logger
func NewLogger(service string) Logger {
	l := logrus.New()
	l.Formatter = new(prettyFormatter)
	return &logger{
		l.WithFields(logrus.Fields{"service": service}),
		make(map[string]*prometheus.HistogramVec),
		make(map[string]prometheus.Counter),
	}
}

// Helper function used to log requests and responses from our GRPC method
// implementations
func (l *logger) Log(request interface{}, response interface{}, err error, duration time.Duration) {
	state := "started"
	if err != nil {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.ErrorLevel, 4)
		state = "errored"
	} else {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.InfoLevel, 4)
	}
	if duration.Seconds() > 0 {
		state = "finished"
	}
	l.ReportMetric(state, duration)
}

func (l *logger) ReportMetric(state string, duration time.Duration) {
	depth := 4
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]
	entry := l.WithFields(logrus.Fields{"method": method})
	fmt.Printf("GOING TO REPORT STATS FOR METHOD: %v\n", method)
	fmt.Printf("state: %v, duration: %v\n", state, duration)

	var tokens []string
	for _, token := range camelcase.Split(method) {
		tokens = append(tokens, strings.ToLower(token))
	}
	rootStatName := strings.Join(tokens, "_")

	runTimeName := fmt.Sprintf("%v_time", rootStatName)
	runTime, ok := l.histogram[runTimeName]
	if !ok {
		bucketFactor := 2.0
		bucketCount := 20 // Which makes the max bucket 2^20 seconds or ~12 days in size
		runTime = prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Namespace: "pachyderm",
				Subsystem: "pachd",
				Name:      runTimeName,
				Help:      fmt.Sprintf("Run time of %v", method),
				Buckets:   prometheus.ExponentialBuckets(1.0, bucketFactor, bucketCount),
			},
			[]string{
				"state", // Since both finished and errored datums can have proc times
			},
		)
		if err := prometheus.Register(runTime); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
		} else {
			l.histogram[runTimeName] = runTime
		}
	}
	if hist, err := runTime.GetMetricWithLabelValues(state); err != nil {
		l.LogAtLevel(entry, logrus.WarnLevel, "failed to get histogram w labels: state (%v) with error %v", state, err)
	} else {
		hist.Observe(duration.Seconds())
	}

	secondsCountName := fmt.Sprintf("%v_seconds_count", rootStatName)
	secondsCount, ok := l.counter[secondsCountName]
	if !ok {
		secondsCount = prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "pachyderm",
				Subsystem: "pachd",
				Name:      secondsCountName,
				Help:      fmt.Sprintf("cumulative number of seconds spent in %v", method),
			},
		)
		if err := prometheus.Register(secondsCount); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
		} else {
			l.counter[secondsCountName] = secondsCount
		}
	}
	secondsCount.Add(duration.Seconds())

}

func (l *logger) LogAtLevel(entry *logrus.Entry, level logrus.Level, args ...interface{}) {
	switch level {
	case logrus.PanicLevel:
		entry.Panic(args)
	case logrus.FatalLevel:
		entry.Fatal(args)
	case logrus.ErrorLevel:
		entry.Error(args)
	case logrus.WarnLevel:
		entry.Warn(args)
	case logrus.InfoLevel:
		entry.Info(args)
	case logrus.DebugLevel:
		entry.Debug(args)
	}
}

func (l *logger) LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int) {
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]

	fields := logrus.Fields{
		"method":  method,
		"request": request,
	}
	if response != nil {
		fields["response"] = response
	}
	if err != nil {
		// "err" itself might be a code or even an empty struct
		fields["error"] = err.Error()
	}
	if duration > 0 {
		fields["duration"] = duration
	}
	l.LogAtLevel(l.WithFields(fields), level)
}

type prettyFormatter struct{}

func (f *prettyFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	serialized := []byte(
		fmt.Sprintf(
			"%v %v ",
			entry.Time.Format(logrus.DefaultTimestampFormat),
			strings.ToUpper(entry.Level.String()),
		),
	)
	if entry.Data["service"] != nil {
		serialized = append(serialized, []byte(fmt.Sprintf("%v.%v ", entry.Data["service"], entry.Data["method"]))...)
	}
	if len(entry.Data) > 2 {
		delete(entry.Data, "service")
		delete(entry.Data, "method")
		if entry.Data["duration"] != nil {
			entry.Data["duration"] = entry.Data["duration"].(time.Duration).Seconds()
		}
		data, err := json.Marshal(entry.Data)
		if err != nil {
			return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
		}
		serialized = append(serialized, []byte(string(data))...)
		serialized = append(serialized, ' ')
	}

	serialized = append(serialized, []byte(entry.Message)...)
	serialized = append(serialized, '\n')
	return serialized, nil
}
<|endoftext|>