<commit_before>/* Data types and structues for concurrent use:
AtomicInt32
AtomicInt64
AtomicBoolean
ConcurrentMap
*/

package tao

import (
	"fmt"
	"sync"
	"sync/atomic"
)

// AtomicInt64 provides atomic int64 type.
type AtomicInt64 int64

// NewAtomicInt64 returns an atomic int64 type.
func NewAtomicInt64(initialValue int64) *AtomicInt64 {
	a := AtomicInt64(initialValue)
	return &a
}

// Get returns the value of int64 atomically.
func (a *AtomicInt64) Get() int64 {
	return int64(*a)
}

// Set sets the value of int64 atomically.
func (a *AtomicInt64) Set(newValue int64) {
	atomic.StoreInt64((*int64)(a), newValue)
}

// GetAndSet sets new value and returns the old atomically.
func (a *AtomicInt64) GetAndSet(newValue int64) int64 {
	for {
		current := a.Get()
		if a.CompareAndSet(current, newValue) {
			return current
		}
	}
}

// CompareAndSet compares int64 with expected value, if equals as expected
// then sets the updated value, this operation performs atomically.
func (a *AtomicInt64) CompareAndSet(expect, update int64) bool {
	return atomic.CompareAndSwapInt64((*int64)(a), expect, update)
}

// GetAndIncrement gets the old value and then increment by 1, this operation
// performs atomically.
func (a *AtomicInt64) GetAndIncrement() int64 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}

}

// GetAndDecrement gets the old value and then decrement by 1, this operation
// performs atomically.
func (a *AtomicInt64) GetAndDecrement() int64 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// GetAndAdd gets the old value and then add by delta, this operation
// performs atomically.
func (a *AtomicInt64) GetAndAdd(delta int64) int64 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// IncrementAndGet increments the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt64) IncrementAndGet() int64 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// DecrementAndGet decrements the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt64) DecrementAndGet() int64 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// AddAndGet adds the value by delta and then gets the value, this operation
// performs atomically.
func (a *AtomicInt64) AddAndGet(delta int64) int64 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

func (a *AtomicInt64) String() string {
	return fmt.Sprintf("%d", a.Get())
}

// AtomicInt32 provides atomic int32 type.
type AtomicInt32 int32

// NewAtomicInt32 returns an atomoic int32 type.
func NewAtomicInt32(initialValue int32) *AtomicInt32 {
	a := AtomicInt32(initialValue)
	return &a
}

// Get returns the value of int32 atomically.
func (a *AtomicInt32) Get() int32 {
	return int32(*a)
}

// Set sets the value of int32 atomically.
func (a *AtomicInt32) Set(newValue int32) {
	atomic.StoreInt32((*int32)(a), newValue)
}

// GetAndSet sets new value and returns the old atomically.
func (a *AtomicInt32) GetAndSet(newValue int32) (oldValue int32) {
	for {
		oldValue = a.Get()
		if a.CompareAndSet(oldValue, newValue) {
			return
		}
	}
}

// CompareAndSet compares int64 with expected value, if equals as expected
// then sets the updated value, this operation performs atomically.
func (a *AtomicInt32) CompareAndSet(expect, update int32) bool {
	return atomic.CompareAndSwapInt32((*int32)(a), expect, update)
}

// GetAndIncrement gets the old value and then increment by 1, this operation
// performs atomically.
func (a *AtomicInt32) GetAndIncrement() int32 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}

}

// GetAndDecrement gets the old value and then decrement by 1, this operation
// performs atomically.
func (a *AtomicInt32) GetAndDecrement() int32 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// GetAndAdd gets the old value and then add by delta, this operation
// performs atomically.
func (a *AtomicInt32) GetAndAdd(delta int32) int32 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// IncrementAndGet increments the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt32) IncrementAndGet() int32 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// DecrementAndGet decrements the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt32) DecrementAndGet() int32 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// AddAndGet adds the value by delta and then gets the value, this operation
// performs atomically.
func (a *AtomicInt32) AddAndGet(delta int32) int32 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

func (a *AtomicInt32) String() string {
	return fmt.Sprintf("%d", a.Get())
}

// AtomicBoolean provides atomic boolean type.
type AtomicBoolean int32

// NewAtomicBoolean returns an atomic boolean type.
func NewAtomicBoolean(initialValue bool) *AtomicBoolean {
	var a AtomicBoolean
	if initialValue {
		a = AtomicBoolean(1)
	} else {
		a = AtomicBoolean(0)
	}
	return &a
}

// Get returns the value of boolean atomically.
func (a *AtomicBoolean) Get() bool {
	return atomic.LoadInt32((*int32)(a)) != 0
}

// Set sets the value of boolean atomically.
func (a *AtomicBoolean) Set(newValue bool) {
	if newValue {
		atomic.StoreInt32((*int32)(a), 1)
	} else {
		atomic.StoreInt32((*int32)(a), 0)
	}
}

// CompareAndSet compares boolean with expected value, if equals as expected
// then sets the updated value, this operation performs atomically.
func (a *AtomicBoolean) CompareAndSet(oldValue, newValue bool) bool {
	var o int32
	var n int32
	if oldValue {
		o = 1
	} else {
		o = 0
	}
	if newValue {
		n = 1
	} else {
		n = 0
	}
	return atomic.CompareAndSwapInt32((*int32)(a), o, n)
}

// GetAndSet sets new value and returns the old atomically.
func (a *AtomicBoolean) GetAndSet(newValue bool) bool {
	for {
		current := a.Get()
		if a.CompareAndSet(current, newValue) {
			return current
		}
	}
}

func (a *AtomicBoolean) String() string {
	return fmt.Sprintf("%t", a.Get())
}

// ConnMap is a safe map for server connection management.
type ConnMap struct {
	sync.RWMutex
	m map[int64]*ServerConn
}

// NewConnMap returns a new ConnMap.
func NewConnMap() *ConnMap {
	return &ConnMap{
		m: make(map[int64]*ServerConn),
	}
}

// Clear clears all elements in map.
func (cm *ConnMap) Clear() {
	cm.Lock()
	cm.m = make(map[int64]*ServerConn)
	cm.Unlock()
}

// Get gets a server connection with specified net ID.
func (cm *ConnMap) Get(id int64) (*ServerConn, bool) {
	cm.RLock()
	sc, ok := cm.m[id]
	cm.RUnlock()
	return sc, ok
}

// Put puts a server connection with specified net ID in map.
func (cm *ConnMap) Put(id int64, sc *ServerConn) {
	cm.Lock()
	cm.m[id] = sc
	cm.Unlock()
}

// Remove removes a server connection with specified net ID.
func (cm *ConnMap) Remove(id int64) {
	cm.Lock()
	delete(cm.m, id)
	cm.Unlock()
}

// Size returns map size.
func (cm *ConnMap) Size() int {
	cm.RLock()
	size := len(cm.m)
	cm.RUnlock()
	return size
}

// IsEmpty tells whether ConnMap is empty.
func (cm *ConnMap) IsEmpty() bool {
	return cm.Size() <= 0
}

// IDs returns all net IDs of ConnMap.
func (cm *ConnMap) IDs() []int64 {
	var ids []int64
	cm.RLock()
	for id := range cm.m {
		ids = append(ids, id)
	}
	cm.Unlock()
	return ids
}
<commit_msg>bugfix unlock of rwmutex<commit_after>/* Data types and structues for concurrent use:
AtomicInt32
AtomicInt64
AtomicBoolean
ConcurrentMap
*/

package tao

import (
	"fmt"
	"sync"
	"sync/atomic"
)

// AtomicInt64 provides atomic int64 type.
type AtomicInt64 int64

// NewAtomicInt64 returns an atomic int64 type.
func NewAtomicInt64(initialValue int64) *AtomicInt64 {
	a := AtomicInt64(initialValue)
	return &a
}

// Get returns the value of int64 atomically.
func (a *AtomicInt64) Get() int64 {
	return int64(*a)
}

// Set sets the value of int64 atomically.
func (a *AtomicInt64) Set(newValue int64) {
	atomic.StoreInt64((*int64)(a), newValue)
}

// GetAndSet sets new value and returns the old atomically.
func (a *AtomicInt64) GetAndSet(newValue int64) int64 {
	for {
		current := a.Get()
		if a.CompareAndSet(current, newValue) {
			return current
		}
	}
}

// CompareAndSet compares int64 with expected value, if equals as expected
// then sets the updated value, this operation performs atomically.
func (a *AtomicInt64) CompareAndSet(expect, update int64) bool {
	return atomic.CompareAndSwapInt64((*int64)(a), expect, update)
}

// GetAndIncrement gets the old value and then increment by 1, this operation
// performs atomically.
func (a *AtomicInt64) GetAndIncrement() int64 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}

}

// GetAndDecrement gets the old value and then decrement by 1, this operation
// performs atomically.
func (a *AtomicInt64) GetAndDecrement() int64 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// GetAndAdd gets the old value and then add by delta, this operation
// performs atomically.
func (a *AtomicInt64) GetAndAdd(delta int64) int64 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// IncrementAndGet increments the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt64) IncrementAndGet() int64 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// DecrementAndGet decrements the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt64) DecrementAndGet() int64 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// AddAndGet adds the value by delta and then gets the value, this operation
// performs atomically.
func (a *AtomicInt64) AddAndGet(delta int64) int64 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

func (a *AtomicInt64) String() string {
	return fmt.Sprintf("%d", a.Get())
}

// AtomicInt32 provides atomic int32 type.
type AtomicInt32 int32

// NewAtomicInt32 returns an atomoic int32 type.
func NewAtomicInt32(initialValue int32) *AtomicInt32 {
	a := AtomicInt32(initialValue)
	return &a
}

// Get returns the value of int32 atomically.
func (a *AtomicInt32) Get() int32 {
	return int32(*a)
}

// Set sets the value of int32 atomically.
func (a *AtomicInt32) Set(newValue int32) {
	atomic.StoreInt32((*int32)(a), newValue)
}

// GetAndSet sets new value and returns the old atomically.
func (a *AtomicInt32) GetAndSet(newValue int32) (oldValue int32) {
	for {
		oldValue = a.Get()
		if a.CompareAndSet(oldValue, newValue) {
			return
		}
	}
}

// CompareAndSet compares int64 with expected value, if equals as expected
// then sets the updated value, this operation performs atomically.
func (a *AtomicInt32) CompareAndSet(expect, update int32) bool {
	return atomic.CompareAndSwapInt32((*int32)(a), expect, update)
}

// GetAndIncrement gets the old value and then increment by 1, this operation
// performs atomically.
func (a *AtomicInt32) GetAndIncrement() int32 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}

}

// GetAndDecrement gets the old value and then decrement by 1, this operation
// performs atomically.
func (a *AtomicInt32) GetAndDecrement() int32 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// GetAndAdd gets the old value and then add by delta, this operation
// performs atomically.
func (a *AtomicInt32) GetAndAdd(delta int32) int32 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return current
		}
	}
}

// IncrementAndGet increments the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt32) IncrementAndGet() int32 {
	for {
		current := a.Get()
		next := current + 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// DecrementAndGet decrements the value by 1 and then gets the value, this
// operation performs atomically.
func (a *AtomicInt32) DecrementAndGet() int32 {
	for {
		current := a.Get()
		next := current - 1
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

// AddAndGet adds the value by delta and then gets the value, this operation
// performs atomically.
func (a *AtomicInt32) AddAndGet(delta int32) int32 {
	for {
		current := a.Get()
		next := current + delta
		if a.CompareAndSet(current, next) {
			return next
		}
	}
}

func (a *AtomicInt32) String() string {
	return fmt.Sprintf("%d", a.Get())
}

// AtomicBoolean provides atomic boolean type.
type AtomicBoolean int32

// NewAtomicBoolean returns an atomic boolean type.
func NewAtomicBoolean(initialValue bool) *AtomicBoolean {
	var a AtomicBoolean
	if initialValue {
		a = AtomicBoolean(1)
	} else {
		a = AtomicBoolean(0)
	}
	return &a
}

// Get returns the value of boolean atomically.
func (a *AtomicBoolean) Get() bool {
	return atomic.LoadInt32((*int32)(a)) != 0
}

// Set sets the value of boolean atomically.
func (a *AtomicBoolean) Set(newValue bool) {
	if newValue {
		atomic.StoreInt32((*int32)(a), 1)
	} else {
		atomic.StoreInt32((*int32)(a), 0)
	}
}

// CompareAndSet compares boolean with expected value, if equals as expected
// then sets the updated value, this operation performs atomically.
func (a *AtomicBoolean) CompareAndSet(oldValue, newValue bool) bool {
	var o int32
	var n int32
	if oldValue {
		o = 1
	} else {
		o = 0
	}
	if newValue {
		n = 1
	} else {
		n = 0
	}
	return atomic.CompareAndSwapInt32((*int32)(a), o, n)
}

// GetAndSet sets new value and returns the old atomically.
func (a *AtomicBoolean) GetAndSet(newValue bool) bool {
	for {
		current := a.Get()
		if a.CompareAndSet(current, newValue) {
			return current
		}
	}
}

func (a *AtomicBoolean) String() string {
	return fmt.Sprintf("%t", a.Get())
}

// ConnMap is a safe map for server connection management.
type ConnMap struct {
	sync.RWMutex
	m map[int64]*ServerConn
}

// NewConnMap returns a new ConnMap.
func NewConnMap() *ConnMap {
	return &ConnMap{
		m: make(map[int64]*ServerConn),
	}
}

// Clear clears all elements in map.
func (cm *ConnMap) Clear() {
	cm.Lock()
	cm.m = make(map[int64]*ServerConn)
	cm.Unlock()
}

// Get gets a server connection with specified net ID.
func (cm *ConnMap) Get(id int64) (*ServerConn, bool) {
	cm.RLock()
	sc, ok := cm.m[id]
	cm.RUnlock()
	return sc, ok
}

// Put puts a server connection with specified net ID in map.
func (cm *ConnMap) Put(id int64, sc *ServerConn) {
	cm.Lock()
	cm.m[id] = sc
	cm.Unlock()
}

// Remove removes a server connection with specified net ID.
func (cm *ConnMap) Remove(id int64) {
	cm.Lock()
	delete(cm.m, id)
	cm.Unlock()
}

// Size returns map size.
func (cm *ConnMap) Size() int {
	cm.RLock()
	size := len(cm.m)
	cm.RUnlock()
	return size
}

// IsEmpty tells whether ConnMap is empty.
func (cm *ConnMap) IsEmpty() bool {
	return cm.Size() <= 0
}

// IDs returns all net IDs of ConnMap.
func (cm *ConnMap) IDs() []int64 {
	var ids []int64
	cm.RLock()
	for id := range cm.m {
		ids = append(ids, id)
	}
	cm.RUnlock()
	return ids
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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 agreedto 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 mysql

import (
	"net"
	"strings"
	"testing"

	"golang.org/x/net/context"
)

func benchmarkQuery(b *testing.B, threads int, query string) {
	th := &testHandler{}

	authServer := &AuthServerNone{}

	l, err := NewListener("tcp", ":0", authServer, th, 0, 0)
	if err != nil {
		b.Fatalf("NewListener failed: %v", err)
	}
	defer l.Close()

	go func() {
		l.Accept()
	}()

	b.SetParallelism(threads)
	b.SetBytes(int64(len(query)))

	host := l.Addr().(*net.TCPAddr).IP.String()
	port := l.Addr().(*net.TCPAddr).Port
	params := &ConnParams{
		Host:  host,
		Port:  port,
		Uname: "user1",
		Pass:  "password1",
	}
	ctx := context.Background()

	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		conn, err := Connect(ctx, params)
		if err != nil {
			b.Fatal(err)
		}
		defer func() {
			conn.writeComQuit()
			conn.Close()
		}()

		for pb.Next() {
			if _, err := conn.ExecuteFetch(query, 1000, true); err != nil {
				b.Fatalf("ExecuteFetch failed: %v", err)
			}
		}
	})
}

// This file contains various long-running tests for mysql.

// BenchmarkParallelShortQueries creates N simultaneous connections, then
// executes M queries on them, then closes them.
// It is meant as a somewhat real load test.
func BenchmarkParallelShortQueries(b *testing.B) {
	benchmarkQuery(b, 10, "select rows")
}

func BenchmarkParallelMediumQueries(b *testing.B) {
	benchmarkQuery(b, 10, "select"+strings.Repeat("x", connBufferSize))
}
<commit_msg>[go/mysql] add benchmark for random query sizes<commit_after>/*
Copyright 2017 Google Inc.

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 agreedto 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 mysql

import (
	"math/rand"
	"net"
	"strings"
	"testing"

	"golang.org/x/net/context"
)

func benchmarkQuery(b *testing.B, threads int, query string) {
	th := &testHandler{}

	authServer := &AuthServerNone{}

	l, err := NewListener("tcp", ":0", authServer, th, 0, 0)
	if err != nil {
		b.Fatalf("NewListener failed: %v", err)
	}
	defer l.Close()

	go func() {
		l.Accept()
	}()

	b.SetParallelism(threads)
	if query != "" {
		b.SetBytes(int64(len(query)))
	}

	host := l.Addr().(*net.TCPAddr).IP.String()
	port := l.Addr().(*net.TCPAddr).Port
	params := &ConnParams{
		Host:  host,
		Port:  port,
		Uname: "user1",
		Pass:  "password1",
	}
	ctx := context.Background()

	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		conn, err := Connect(ctx, params)
		if err != nil {
			b.Fatal(err)
		}
		defer func() {
			conn.writeComQuit()
			conn.Close()
		}()

		for pb.Next() {
			execQuery := query
			if execQuery == "" {
				// generate random query
				n := rand.Intn(MaxPacketSize-2) + 1
				execQuery = strings.Repeat("x", n)
			}
			if _, err := conn.ExecuteFetch(execQuery, 1000, true); err != nil {
				b.Fatalf("ExecuteFetch failed: %v", err)
			}
		}
	})
}

// This file contains various long-running tests for mysql.

// BenchmarkParallelShortQueries creates N simultaneous connections, then
// executes M queries on them, then closes them.
// It is meant as a somewhat real load test.
func BenchmarkParallelShortQueries(b *testing.B) {
	benchmarkQuery(b, 10, "select rows")
}

func BenchmarkParallelMediumQueries(b *testing.B) {
	benchmarkQuery(b, 10, "select"+strings.Repeat("x", connBufferSize))
}

func BenchmarkParallelRandomQueries(b *testing.B) {
	benchmarkQuery(b, 10, "")
}
<|endoftext|>


<commit_before>package pprint

import (
	"io"
	"strings"
)

// toStream recursively converts a document into the stream elements
// we'll be using.  We use channels to organize the coroutines.
func toStream(document Element) <-chan streamElt {
	ch := make(chan streamElt)
	go func() {
		visitElement(document, ch)
	}()
	return ch
}

func visitElement(document Element, out chan<- streamElt) {
	switch doc := document.(type) {
	case *Text:
		out <- &textElt{elt{-1}, doc.text}
	case *Cond:
		out <- &condElt{elt{-1}, doc.small, doc.continuation, doc.tail}
	case *LineBreak:
		out <- &crlfElt{elt{-1}}
	case *Concat:
		for _, elt := range doc.children {
			visitElement(elt, out)
		}
	case *Group:
		out <- &gbegElt{elt{-1}}
		visitElement(doc.child, out)
		out <- &gendElt{elt{-1}}
	case *Nest:
		out <- &nbegElt{elt{-1}}
		out <- &gbegElt{elt{-1}}
		visitElement(doc.child, out)
		out <- &gendElt{elt{-1}}
		out <- &nendElt{elt{-1}}
	default:
		panic("Couldn't understand document type")
	}
}

// annotateLastChar is the next step; it takes the stream elements
// from `toStream` and adds information about the horizontal position
// of their last character.  This is not possible with NBeg and GBeg
// elements as we haven't got enough information yet.
func annotateLastChar(in <-chan streamElt) <-chan streamElt {
	ch := make(chan streamElt)
	go func() {
		position := 0
		for {
			select {
			case elt, ok := <-in:
				if !ok {
					return
				}
				switch elt := elt.(type) {
				case *textElt:
					position += len(elt.payload)
					elt.hpos = position
					ch <- elt
				case *condElt:
					position += len(elt.small)
					elt.hpos = position
					ch <- elt
				case *crlfElt:
					elt.hpos = position
					ch <- elt
				case *gbegElt, *nbegElt:
					// Don't have enough information yet to do this
					// accurately.
					ch <- elt
				case *gendElt:
					elt.hpos = position
					ch <- elt
				case *nendElt:
					elt.hpos = position
					ch <- elt
				}
			}
		}
	}()
	return ch
}

type lookaheadStack []streamElt

// annotateGBeg is the next step; we take the horizontal position
// information gotten from `annotateLastChar` and compute the `hpos`
// for GBeg elements.  We don't need to do it for NBeg, but for GBeg
// it matters for linebreaks.
func annotateGBeg(in <-chan streamElt) <-chan streamElt {
	ch := make(chan streamElt)
	go func() {
		var lookahead []lookaheadStack
		for {
			select {
			case element, ok := <-in:
				if !ok {
					return
				}
				switch element := element.(type) {
				case *textElt, *condElt, *crlfElt, *nbegElt, *nendElt:
					if len(lookahead) == 0 {
						ch <- element
					} else {
						last := len(lookahead) - 1
						lookahead[last] = append(lookahead[last], element)
					}
				case *gbegElt:
					newList := make(lookaheadStack, 0)
					lookahead = append(lookahead, newList)
				case *gendElt:
					last := len(lookahead) - 1
					top := lookahead[last]
					lookahead = lookahead[0:last]
					if len(lookahead) == 0 {
						// this, then, was the topmost stack
						ch <- &gbegElt{elt{element.hpos}}
						for _, e := range top {
							ch <- e
						}
						ch <- element
					} else {
						newtop := lookahead[last-1]
						newtop = append(newtop, &gbegElt{elt{element.hpos}})
						newtop = append(newtop, top...)
						newtop = append(newtop, element)
						lookahead[last-1] = newtop
					}
				}
			}
		}
	}()
	return ch
}

// Kiselyov's original formulation includes an alternate third phase
// which limits lookahead to the width of the page.  This is difficult
// for us because we don't guarantee docs are of nonzero length,
// although that could be finessed, and also it adds extra complexity
// for minimal benefit.  This implementation skips it.

// The final phase is to compute output.  Each time we see a
// `gbeg_element_t`, we can compare its `hpos` with `rightEdge` to see
// whether it'll fit without breaking.  If it does fit, increment
// `fittingElements` and proceed, which will cause the logic for
// `text_element_t` and `cond_element_t` to just append stuff without
// line breaks.  If it doesn't fit, set `fittingElements` to 0, which
// will cause `cond_element_t` to do line breaks.  When we do a line
// break, we need to compute where the new right edge of the 'page'
// would be in the context of the original stream; so if we saw a
// `cond_element_t` with `e.hpos` of 300 (meaning it ends at
// horizontal position 300), the new right edge would be 300 -
// indentation + page width.
func output(in <-chan streamElt, width int, output io.Writer) error {
	fittingElements := 0
	rightEdge := width
	hpos := 0
	var indent []int
	for {
		select {
		case elt, ok := <-in:
			if !ok {
				return nil
			}
			switch elt := elt.(type) {
			case *textElt:
				_, err := io.WriteString(output, elt.payload)
				if err != nil {
					return err
				}
				hpos += len(elt.payload)
			case *condElt:
				if fittingElements == 0 {
					var currentIndent int
					if len(indent) == 0 {
						currentIndent = 0
					} else {
						currentIndent = indent[len(indent)-1]
					}
					_, err := io.WriteString(output, elt.tail)
					if err != nil {
						return err
					}
					_, err = io.WriteString(output, "\n")
					if err != nil {
						return err
					}
					_, err = io.WriteString(output, strings.Repeat(" ", currentIndent))
					if err != nil {
						return err
					}
					_, err = io.WriteString(output, elt.cont)
					if err != nil {
						return err
					}
					fittingElements = 0
					hpos = currentIndent + len(elt.cont)
					rightEdge = (width - hpos) + elt.hpos
				} else {
					_, err := io.WriteString(output, elt.small)
					if err != nil {
						return err
					}
					hpos += len(elt.small)
				}
			case *crlfElt:
				var currentIndent int
				if len(indent) == 0 {
					currentIndent = 0
				} else {
					currentIndent = indent[len(indent)-1]
				}
				_, err := io.WriteString(output, "\n")
				if err != nil {
					return err
				}
				_, err = io.WriteString(output, strings.Repeat(" ", currentIndent))
				if err != nil {
					return err
				}
				fittingElements = 0
				hpos = currentIndent
				rightEdge = (width - hpos) + elt.hpos
			case *gbegElt:
				if fittingElements != 0 || elt.hpos <= rightEdge {
					fittingElements++
				} else {
					fittingElements = 0
				}
			case *gendElt:
				if fittingElements != 0 {
					fittingElements--
				}
			case *nbegElt:
				indent = append(indent, hpos)
			case *nendElt:
				if len(indent) > 0 {
					indent = indent[0 : len(indent)-1]
				}
			}
		}
	}
}

// PrettyPrint prints `doc` to `out` assuming a right page edge of
// `width`.
func PrettyPrint(doc Element, width int, out io.Writer) error {
	return output(annotateGBeg(annotateLastChar(toStream(doc))), width, out)
}
<commit_msg>Actually close things, derp.<commit_after>package pprint

import (
	"io"
	"strings"
)

// toStream recursively converts a document into the stream elements
// we'll be using.  We use channels to organize the coroutines.
func toStream(document Element) <-chan streamElt {
	ch := make(chan streamElt)
	go func() {
		defer close(ch)
		visitElement(document, ch)
	}()
	return ch
}

func visitElement(document Element, out chan<- streamElt) {
	switch doc := document.(type) {
	case *Text:
		out <- &textElt{elt{-1}, doc.text}
	case *Cond:
		out <- &condElt{elt{-1}, doc.small, doc.continuation, doc.tail}
	case *LineBreak:
		out <- &crlfElt{elt{-1}}
	case *Concat:
		for _, elt := range doc.children {
			visitElement(elt, out)
		}
	case *Group:
		out <- &gbegElt{elt{-1}}
		visitElement(doc.child, out)
		out <- &gendElt{elt{-1}}
	case *Nest:
		out <- &nbegElt{elt{-1}}
		out <- &gbegElt{elt{-1}}
		visitElement(doc.child, out)
		out <- &gendElt{elt{-1}}
		out <- &nendElt{elt{-1}}
	default:
		panic("Couldn't understand document type")
	}
}

// annotateLastChar is the next step; it takes the stream elements
// from `toStream` and adds information about the horizontal position
// of their last character.  This is not possible with NBeg and GBeg
// elements as we haven't got enough information yet.
func annotateLastChar(in <-chan streamElt) <-chan streamElt {
	ch := make(chan streamElt)
	go func() {
		defer close(ch)
		position := 0
		for {
			select {
			case elt, ok := <-in:
				if !ok {
					return
				}
				switch elt := elt.(type) {
				case *textElt:
					position += len(elt.payload)
					elt.hpos = position
					ch <- elt
				case *condElt:
					position += len(elt.small)
					elt.hpos = position
					ch <- elt
				case *crlfElt:
					elt.hpos = position
					ch <- elt
				case *gbegElt, *nbegElt:
					// Don't have enough information yet to do this
					// accurately.
					ch <- elt
				case *gendElt:
					elt.hpos = position
					ch <- elt
				case *nendElt:
					elt.hpos = position
					ch <- elt
				}
			}
		}
	}()
	return ch
}

type lookaheadStack []streamElt

// annotateGBeg is the next step; we take the horizontal position
// information gotten from `annotateLastChar` and compute the `hpos`
// for GBeg elements.  We don't need to do it for NBeg, but for GBeg
// it matters for linebreaks.
func annotateGBeg(in <-chan streamElt) <-chan streamElt {
	ch := make(chan streamElt)
	go func() {
		defer close(ch)
		var lookahead []lookaheadStack
		for {
			select {
			case element, ok := <-in:
				if !ok {
					return
				}
				switch element := element.(type) {
				case *textElt, *condElt, *crlfElt, *nbegElt, *nendElt:
					if len(lookahead) == 0 {
						ch <- element
					} else {
						last := len(lookahead) - 1
						lookahead[last] = append(lookahead[last], element)
					}
				case *gbegElt:
					newList := make(lookaheadStack, 0)
					lookahead = append(lookahead, newList)
				case *gendElt:
					last := len(lookahead) - 1
					top := lookahead[last]
					lookahead = lookahead[0:last]
					if len(lookahead) == 0 {
						// this, then, was the topmost stack
						ch <- &gbegElt{elt{element.hpos}}
						for _, e := range top {
							ch <- e
						}
						ch <- element
					} else {
						newtop := lookahead[last-1]
						newtop = append(newtop, &gbegElt{elt{element.hpos}})
						newtop = append(newtop, top...)
						newtop = append(newtop, element)
						lookahead[last-1] = newtop
					}
				}
			}
		}
	}()
	return ch
}

// Kiselyov's original formulation includes an alternate third phase
// which limits lookahead to the width of the page.  This is difficult
// for us because we don't guarantee docs are of nonzero length,
// although that could be finessed, and also it adds extra complexity
// for minimal benefit.  This implementation skips it.

// The final phase is to compute output.  Each time we see a
// `gbeg_element_t`, we can compare its `hpos` with `rightEdge` to see
// whether it'll fit without breaking.  If it does fit, increment
// `fittingElements` and proceed, which will cause the logic for
// `text_element_t` and `cond_element_t` to just append stuff without
// line breaks.  If it doesn't fit, set `fittingElements` to 0, which
// will cause `cond_element_t` to do line breaks.  When we do a line
// break, we need to compute where the new right edge of the 'page'
// would be in the context of the original stream; so if we saw a
// `cond_element_t` with `e.hpos` of 300 (meaning it ends at
// horizontal position 300), the new right edge would be 300 -
// indentation + page width.
func output(in <-chan streamElt, width int, output io.Writer) error {
	fittingElements := 0
	rightEdge := width
	hpos := 0
	var indent []int
	for {
		select {
		case elt, ok := <-in:
			if !ok {
				return nil
			}
			switch elt := elt.(type) {
			case *textElt:
				_, err := io.WriteString(output, elt.payload)
				if err != nil {
					return err
				}
				hpos += len(elt.payload)
			case *condElt:
				if fittingElements == 0 {
					var currentIndent int
					if len(indent) == 0 {
						currentIndent = 0
					} else {
						currentIndent = indent[len(indent)-1]
					}
					_, err := io.WriteString(output, elt.tail)
					if err != nil {
						return err
					}
					_, err = io.WriteString(output, "\n")
					if err != nil {
						return err
					}
					_, err = io.WriteString(output, strings.Repeat(" ", currentIndent))
					if err != nil {
						return err
					}
					_, err = io.WriteString(output, elt.cont)
					if err != nil {
						return err
					}
					fittingElements = 0
					hpos = currentIndent + len(elt.cont)
					rightEdge = (width - hpos) + elt.hpos
				} else {
					_, err := io.WriteString(output, elt.small)
					if err != nil {
						return err
					}
					hpos += len(elt.small)
				}
			case *crlfElt:
				var currentIndent int
				if len(indent) == 0 {
					currentIndent = 0
				} else {
					currentIndent = indent[len(indent)-1]
				}
				_, err := io.WriteString(output, "\n")
				if err != nil {
					return err
				}
				_, err = io.WriteString(output, strings.Repeat(" ", currentIndent))
				if err != nil {
					return err
				}
				fittingElements = 0
				hpos = currentIndent
				rightEdge = (width - hpos) + elt.hpos
			case *gbegElt:
				if fittingElements != 0 || elt.hpos <= rightEdge {
					fittingElements++
				} else {
					fittingElements = 0
				}
			case *gendElt:
				if fittingElements != 0 {
					fittingElements--
				}
			case *nbegElt:
				indent = append(indent, hpos)
			case *nendElt:
				if len(indent) > 0 {
					indent = indent[0 : len(indent)-1]
				}
			}
		}
	}
}

// PrettyPrint prints `doc` to `out` assuming a right page edge of
// `width`.
func PrettyPrint(doc Element, width int, out io.Writer) error {
	return output(annotateGBeg(annotateLastChar(toStream(doc))), width, out)
}
<|endoftext|>


<commit_before>// Slog is an abbrevation for simplelog. It uses by default stdout as the output
// destination and has features like defining multiple output destinations
// (like stdout and file at the same time), prefix pretending with a function call
// and global switch to turn off/on logs.
package slog

import (
	"errors"
	"fmt"
	"io"
	"os"
	"sync"
	"time"
)

// A slog represents an object that generates lines of output to to an
// io.Writer. By default it uses os.Stdout, but it can be changed  or others
// may be included during creation.
type Slog struct {
	mu      sync.Mutex    // protects the following fields
	disable bool          // global switch to disable log completely
	prefix  func() string // function return is written at beginning of each line
	out     io.Writer     // destination for ouput
}

var stdlog = New()

// New creates a new slog. The filepath sets the files that will be used
// as an extra output destination. By default slog outputs to stdout.
func New(filepath ...string) *Slog {
	writers := make([]io.Writer, 0)
	for _, path := range filepath {
		logFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0640)
		if err != nil {
			fmt.Printf("slog: can't open %s: '%s'\n", path, err)
			continue
		}

		writers = append(writers, logFile)
	}

	writers = append(writers, os.Stdout)

	return &Slog{
		out: io.MultiWriter(writers...),
		prefix: func() string {
			return fmt.Sprintf("[%s] ", time.Now().Format(time.StampMilli))
		},
	}
}

// Print calls Output to print to the standard logger. Arguments are handled in
// the manner of fmt.Print.
func Printn(v ...interface{}) (int, error) {
	if stdlog.disable {
		return 0, nil
	}

	return fmt.Fprint(stdlog.output(), v...)
}

// Printf calls Output to print to the standard logger. Arguments are handled in
// the manner of fmt.Printf.
func Printf(format string, v ...interface{}) (int, error) {
	if stdlog.disable {
		return 0, nil
	}

	return fmt.Fprintf(stdlog.output(), format, v...)
}

// Println calls Output to print to the standard logger. Arguments are handled in
// the manner of fmt.Println.
func Println(v ...interface{}) (int, error) {
	if stdlog.disable {
		return 0, nil
	}

	return fmt.Fprintln(stdlog.output(), v...)
}

// SetPrefix sets the output prefix according to the return value of the passed
// function for the standard logger.
func SetPrefix(fn func() string) {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.prefix = fn
}

// Prefix returns the output prefix for the standard logger.
func Prefix() string {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	return stdlog.prefix()
}

// DisablePrefix disables the prefix generator. That means it will reset the current
// prefix generation function.
func DisablePrefix() {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.prefix = func() string { return "" }
}

// SetName adds an additional prefix to the beginning of each line. Useful to
// add an application name. This will modify the the default Prefix() function.
func SetName(name string) {
	fn := stdlog.prefix() // copy the function to prevent recursion
	stdlog.prefix = func() string { return fmt.Sprintf("%s %s", name, fn) }
}

// SetOutput replaces the standard destination for the standard logger.
func SetOutput(out io.Writer) {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.out = out
}

// SetOutputFile defines a new file destination for the output. That means
// if used, it will store the stdout to an external file too.
func SetOutputFile(path string) error {
	if path == "" {
		return errors.New("slog: arg for SetOutputFile sould be not empty")
	}

	writers := make([]io.Writer, 0)
	logFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return fmt.Errorf("slog: can't open %s: '%s'\n", path, err)
	} else {
		stdlog.mu.Lock()
		defer stdlog.mu.Unlock()

		writers = append(writers, logFile, os.Stdout)
		stdlog.out = io.MultiWriter(writers...)
	}

	return nil
}

// DisableLog is a global switch that disables the output completely for the
// standard logger. Useful if you want turn off/on logs for debugging.
func DisableLog() {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.disable = true
}

// Print formats using the default formats for its operands and writes to
// standard output. Spaces are added between operands when neither is a string. It
// returns the number of bytes written and any write error encountered.
func (s *Slog) Printn(v ...interface{}) (int, error) {
	if s.disable {
		return 0, nil
	}

	return fmt.Fprint(s.output(), v...)
}

// Printf formats according to a format specifier and writes to standard output.
// It returns the number of bytes written and any write error encountered.
func (s *Slog) Printf(format string, v ...interface{}) (int, error) {
	if s.disable {
		return 0, nil
	}

	return fmt.Fprintf(s.output(), format, v...)
}

// Println formats using the default formats for its operands and writes to
// standard output. Spaces are always added between operands and a newline is
// appended. It returns the number of bytes written and any write error
// encountered.
func (s *Slog) Println(v ...interface{}) (int, error) {
	if s.disable {
		return 0, nil
	}

	return fmt.Fprintln(s.output(), v...)
}

// SetPrefix sets the output prefix according to the return value of the passed
// function.
func (s *Slog) SetPrefix(fn func() string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.prefix = fn
}

// Prefix returns the output prefix.
func (s *Slog) Prefix() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.prefix()
}

func (s *Slog) DisablePrefix() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.prefix = func() string { return "" }
}

// SetName adds an additional prefix to the beginning of each line. Useful to
// add an application name. This will modify the the default Prefix() function.
func (s *Slog) SetName(name string) {
	fn := s.prefix() // copy the function to prevent recursion
	s.prefix = func() string { return fmt.Sprintf("%s %s", name, fn) }
}

// SetOutput replaces the standard destination.
func (s *Slog) SetOutput(out io.Writer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.out = out
}

// DisableLog is a global switch that disables the output completely. Useful
// if you want turn off/on logs for debugging.
func (s *Slog) DisableLog() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.disable = true
}

func (s *Slog) output() io.Writer {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.out.Write([]byte(s.prefix()))
	return s.out
}
<commit_msg>Add timeLayout, prefixFunc and many other improvements<commit_after>// Slog is an abbrevation for simplelog. It uses by default stdout as the output
// destination and has features like defining multiple output destinations
// (like stdout and file at the same time), prefix pretending with a function call
// and global switch to turn off/on logs.
package slog

import (
	"errors"
	"fmt"
	"io"
	"os"
	"sync"
	"time"
)

// A slog represents an object that generates lines of output to to an
// io.Writer. By default it uses os.Stdout, but it can be changed  or others
// may be included during creation.
type Slog struct {
	mu         sync.Mutex    // protects the following fields
	disable    bool          // global switch to disable log completely
	prefixFunc func() string // if set function return is used instead of prefix
	prefix     struct {
		name       string // is written at beginning of each line
		timeLayout string // is written after prefix.name, used for time.Format()
	}
	out io.Writer // destination for ouput
}

var stdlog = New()

// New creates a new slog. The filepath sets the files that will be used
// as an extra output destination. By default slog outputs to stdout.
func New(filepath ...string) *Slog {
	writers := make([]io.Writer, 0)
	for _, path := range filepath {
		logFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0640)
		if err != nil {
			fmt.Printf("slog: can't open %s: '%s'\n", path, err)
			continue
		}

		writers = append(writers, logFile)
	}

	writers = append(writers, os.Stdout)

	s := &Slog{
		out: io.MultiWriter(writers...),
	}
	s.prefix.timeLayout = time.StampMilli
	return s
}

// Print calls Output to print to the standard logger. Arguments are handled in
// the manner of fmt.Print.
func Printn(v ...interface{}) (int, error) {
	if stdlog.disable {
		return 0, nil
	}

	return fmt.Fprint(stdlog.output(), v...)
}

// Printf calls Output to print to the standard logger. Arguments are handled in
// the manner of fmt.Printf.
func Printf(format string, v ...interface{}) (int, error) {
	if stdlog.disable {
		return 0, nil
	}

	return fmt.Fprintf(stdlog.output(), format, v...)
}

// Println calls Output to print to the standard logger. Arguments are handled in
// the manner of fmt.Println.
func Println(v ...interface{}) (int, error) {
	if stdlog.disable {
		return 0, nil
	}

	return fmt.Fprintln(stdlog.output(), v...)
}

// SetPrefixFunc sets the output prefix according to the return value of the passed
// function for the standard logger.
func SetPrefixFunc(fn func() string) {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.prefixFunc = fn
}

// Prefix returns the output prefix for the standard logger.
func Prefix() string {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	return stdlog.prefixOutput()
}

// DisablePrefix disables the prefix generator for the standard logger. That
// means it will reset the current prefix generation function.
func DisablePrefix() {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.prefixFunc = func() string { return "" }
}

// SetName adds an additional prefix to the beginning of each line for the
// standard logger. Useful to add an application name. This will modify the the
// default Prefix() function.
func SetName(name string) {
	stdlog.prefix.name = name
}

// SetTimeStamp changes the timestamp that is generated in the prefix generator
// function for the standard logger. It's passed to time.Format(). Example
// layouts are: time.RFC3339, time.Kitchen, time.UnixDate ...
func SetTimeStamp(layout string) {
	stdlog.prefix.timeLayout = layout
}

// SetOutput replaces the standard destination for the standard logger.
func SetOutput(out io.Writer) {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.out = out
}

// SetOutputFile defines a new file destination for the output. That means
// if used, it will store the stdout to an external file too.
func SetOutputFile(path string) error {
	if path == "" {
		return errors.New("slog: arg for SetOutputFile sould be not empty")
	}

	writers := make([]io.Writer, 0)
	logFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return fmt.Errorf("slog: can't open %s: '%s'\n", path, err)
	} else {
		stdlog.mu.Lock()
		defer stdlog.mu.Unlock()

		writers = append(writers, logFile, os.Stdout)
		stdlog.out = io.MultiWriter(writers...)
	}

	return nil
}

// DisableLog is a global switch that disables the output completely for the
// standard logger. Useful if you want turn off/on logs for debugging.
func DisableLog() {
	stdlog.mu.Lock()
	defer stdlog.mu.Unlock()
	stdlog.disable = true
}

// Print formats using the default formats for its operands and writes to
// standard output. Spaces are added between operands when neither is a string. It
// returns the number of bytes written and any write error encountered.
func (s *Slog) Printn(v ...interface{}) (int, error) {
	if s.disable {
		return 0, nil
	}

	return fmt.Fprint(s.output(), v...)
}

// Printf formats according to a format specifier and writes to standard output.
// It returns the number of bytes written and any write error encountered.
func (s *Slog) Printf(format string, v ...interface{}) (int, error) {
	if s.disable {
		return 0, nil
	}

	return fmt.Fprintf(s.output(), format, v...)
}

// Println formats using the default formats for its operands and writes to
// standard output. Spaces are always added between operands and a newline is
// appended. It returns the number of bytes written and any write error
// encountered.
func (s *Slog) Println(v ...interface{}) (int, error) {
	if s.disable {
		return 0, nil
	}

	return fmt.Fprintln(s.output(), v...)
}

// SetPrefixFunc sets the output prefix according to the return value of the passed
// function.
func (s *Slog) SetPrefixFunc(fn func() string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.prefixFunc = fn
}

// Prefix returns the output prefix.
func (s *Slog) Prefix() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.prefixOutput()
}

// DisablePrefix disables the prefix generator. That means it will reset the current
// prefix generation function.
func (s *Slog) DisablePrefix() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.prefixFunc = func() string { return "" }
}

// SetName adds an additional prefix to the beginning of each line. Useful to
// add an application name. This will modify the the default Prefix() function.
func (s *Slog) SetName(name string) {
	s.prefix.name = name
}

// SetTimeStamp changes the timestamp that is generated in the prefix generator
// function. It's passed to time.Format(). Example layouts are: time.RFC3339,
// time.Kitchen, time.UnixDate ...
func (s *Slog) SetTimeStamp(layout string) {
	s.prefix.timeLayout = layout
}

// SetOutput replaces the standard destination.
func (s *Slog) SetOutput(out io.Writer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.out = out
}

// DisableLog is a global switch that disables the output completely. Useful
// if you want turn off/on logs for debugging.
func (s *Slog) DisableLog() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.disable = true
}

func (s *Slog) output() io.Writer {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.out.Write([]byte(s.prefixOutput()))
	return s.out
}

func (s *Slog) prefixOutput() string {
	if s.prefixFunc != nil {
		return s.prefixFunc()
	}

	if s.prefix.name == "" {
		return fmt.Sprintf("[%s] ", time.Now().Format(s.prefix.timeLayout))
	}
	return fmt.Sprintf("[%s %s] ", s.prefix.name, time.Now().Format(s.prefix.timeLayout))
}
<|endoftext|>


<commit_before>// Inferno utils/6l/obj.c
// https://bitbucket.org/inferno-os/inferno-os/src/default/utils/6l/obj.c
//
//	Copyright © 1994-1999 Lucent Technologies Inc.  All rights reserved.
//	Portions Copyright © 1995-1997 C H Forsyth (forsyth@terzarima.net)
//	Portions Copyright © 1997-1999 Vita Nuova Limited
//	Portions Copyright © 2000-2007 Vita Nuova Holdings Limited (www.vitanuova.com)
//	Portions Copyright © 2004,2006 Bruce Ellis
//	Portions Copyright © 2005-2007 C H Forsyth (forsyth@terzarima.net)
//	Revisions Copyright © 2000-2007 Lucent Technologies Inc. and others
//	Portions Copyright © 2009 The Go Authors. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package ld

import (
	"bufio"
	"cmd/internal/objabi"
	"cmd/internal/sys"
	"cmd/link/internal/sym"
	"flag"
	"log"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
)

var (
	pkglistfornote []byte
	windowsgui     bool // writes a "GUI binary" instead of a "console binary"
	ownTmpDir      bool // set to true if tmp dir created by linker (e.g. no -tmpdir)
)

func init() {
	flag.Var(&rpath, "r", "set the ELF dynamic linker search `path` to dir1:dir2:...")
}

// Flags used by the linker. The exported flags are used by the architecture-specific packages.
var (
	flagBuildid = flag.String("buildid", "", "record `id` as Go toolchain build id")

	flagOutfile    = flag.String("o", "", "write output to `file`")
	flagPluginPath = flag.String("pluginpath", "", "full path name for plugin")

	flagInstallSuffix = flag.String("installsuffix", "", "set package directory `suffix`")
	flagDumpDep       = flag.Bool("dumpdep", false, "dump symbol dependency graph")
	flagRace          = flag.Bool("race", false, "enable race detector")
	flagMsan          = flag.Bool("msan", false, "enable MSan interface")

	flagFieldTrack = flag.String("k", "", "set field tracking `symbol`")
	flagLibGCC     = flag.String("libgcc", "", "compiler support lib for internal linking; use \"none\" to disable")
	flagTmpdir     = flag.String("tmpdir", "", "use `directory` for temporary files")

	flagExtld      = flag.String("extld", "", "use `linker` when linking in external mode")
	flagExtldflags = flag.String("extldflags", "", "pass `flags` to external linker")
	flagExtar      = flag.String("extar", "", "archive program for buildmode=c-archive")

	flagA           = flag.Bool("a", false, "disassemble output")
	FlagC           = flag.Bool("c", false, "dump call graph")
	FlagD           = flag.Bool("d", false, "disable dynamic executable")
	flagF           = flag.Bool("f", false, "ignore version mismatch")
	flagG           = flag.Bool("g", false, "disable go package data checks")
	flagH           = flag.Bool("h", false, "halt on error")
	flagN           = flag.Bool("n", false, "dump symbol table")
	FlagS           = flag.Bool("s", false, "disable symbol table")
	flagU           = flag.Bool("u", false, "reject unsafe packages")
	FlagW           = flag.Bool("w", false, "disable DWARF generation")
	Flag8           bool // use 64-bit addresses in symbol table
	flagInterpreter = flag.String("I", "", "use `linker` as ELF dynamic linker")
	FlagDebugTramp  = flag.Int("debugtramp", 0, "debug trampolines")
	FlagStrictDups  = flag.Int("strictdups", 0, "sanity check duplicate symbol contents during object file reading (1=warn 2=err).")
	FlagNewLdElf    = flag.Bool("newldelf", false, "ELF host obj load with new loader")

	FlagRound       = flag.Int("R", -1, "set address rounding `quantum`")
	FlagTextAddr    = flag.Int64("T", -1, "set text segment `address`")
	flagEntrySymbol = flag.String("E", "", "set `entry` symbol name")

	cpuprofile     = flag.String("cpuprofile", "", "write cpu profile to `file`")
	memprofile     = flag.String("memprofile", "", "write memory profile to `file`")
	memprofilerate = flag.Int64("memprofilerate", 0, "set runtime.MemProfileRate to `rate`")
)

// Main is the main entry point for the linker code.
func Main(arch *sys.Arch, theArch Arch) {
	thearch = theArch
	ctxt := linknew(arch)
	ctxt.Bso = bufio.NewWriter(os.Stdout)

	// For testing behavior of go command when tools crash silently.
	// Undocumented, not in standard flag parser to avoid
	// exposing in usage message.
	for _, arg := range os.Args {
		if arg == "-crash_for_testing" {
			os.Exit(2)
		}
	}

	final := gorootFinal()
	addstrdata1(ctxt, "runtime/internal/sys.DefaultGoroot="+final)
	addstrdata1(ctxt, "cmd/internal/objabi.defaultGOROOT="+final)

	// TODO(matloob): define these above and then check flag values here
	if ctxt.Arch.Family == sys.AMD64 && objabi.GOOS == "plan9" {
		flag.BoolVar(&Flag8, "8", false, "use 64-bit addresses in symbol table")
	}
	flagHeadType := flag.String("H", "", "set header `type`")
	flag.BoolVar(&ctxt.linkShared, "linkshared", false, "link against installed Go shared libraries")
	flag.Var(&ctxt.LinkMode, "linkmode", "set link `mode`")
	flag.Var(&ctxt.BuildMode, "buildmode", "set build `mode`")
	flag.BoolVar(&ctxt.compressDWARF, "compressdwarf", true, "compress DWARF if possible")
	objabi.Flagfn1("B", "add an ELF NT_GNU_BUILD_ID `note` when using ELF", addbuildinfo)
	objabi.Flagfn1("L", "add specified `directory` to library path", func(a string) { Lflag(ctxt, a) })
	objabi.AddVersionFlag() // -V
	objabi.Flagfn1("X", "add string value `definition` of the form importpath.name=value", func(s string) { addstrdata1(ctxt, s) })
	objabi.Flagcount("v", "print link trace", &ctxt.Debugvlog)
	objabi.Flagfn1("importcfg", "read import configuration from `file`", ctxt.readImportCfg)

	objabi.Flagparse(usage)

	switch *flagHeadType {
	case "":
	case "windowsgui":
		ctxt.HeadType = objabi.Hwindows
		windowsgui = true
	default:
		if err := ctxt.HeadType.Set(*flagHeadType); err != nil {
			Errorf(nil, "%v", err)
			usage()
		}
	}

	if objabi.Fieldtrack_enabled != 0 {
		ctxt.Reachparent = make(map[*sym.Symbol]*sym.Symbol)
	}
	checkStrictDups = *FlagStrictDups

	startProfile()
	if ctxt.BuildMode == BuildModeUnset {
		ctxt.BuildMode = BuildModeExe
	}

	if ctxt.BuildMode != BuildModeShared && flag.NArg() != 1 {
		usage()
	}

	if *flagOutfile == "" {
		*flagOutfile = "a.out"
		if ctxt.HeadType == objabi.Hwindows {
			*flagOutfile += ".exe"
		}
	}

	interpreter = *flagInterpreter

	libinit(ctxt) // creates outfile

	if ctxt.HeadType == objabi.Hunknown {
		ctxt.HeadType.Set(objabi.GOOS)
	}

	ctxt.computeTLSOffset()
	thearch.Archinit(ctxt)

	if ctxt.linkShared && !ctxt.IsELF {
		Exitf("-linkshared can only be used on elf systems")
	}

	if ctxt.Debugvlog != 0 {
		ctxt.Logf("HEADER = -H%d -T0x%x -R0x%x\n", ctxt.HeadType, uint64(*FlagTextAddr), uint32(*FlagRound))
	}

	switch ctxt.BuildMode {
	case BuildModeShared:
		for i := 0; i < flag.NArg(); i++ {
			arg := flag.Arg(i)
			parts := strings.SplitN(arg, "=", 2)
			var pkgpath, file string
			if len(parts) == 1 {
				pkgpath, file = "main", arg
			} else {
				pkgpath, file = parts[0], parts[1]
			}
			pkglistfornote = append(pkglistfornote, pkgpath...)
			pkglistfornote = append(pkglistfornote, '\n')
			addlibpath(ctxt, "command line", "command line", file, pkgpath, "")
		}
	case BuildModePlugin:
		addlibpath(ctxt, "command line", "command line", flag.Arg(0), *flagPluginPath, "")
	default:
		addlibpath(ctxt, "command line", "command line", flag.Arg(0), "main", "")
	}
	ctxt.loadlib()

	deadcode(ctxt)

	if *FlagNewLdElf {
		ctxt.linksetup()
	}

	ctxt.loadlibfull() // XXX do it here for now

	if !*FlagNewLdElf {
		ctxt.linksetupold()
	}
	ctxt.dostrdata()
	dwarfGenerateDebugInfo(ctxt)

	if objabi.Fieldtrack_enabled != 0 {
		fieldtrack(ctxt)
	}
	ctxt.mangleTypeSym()
	ctxt.callgraph()

	ctxt.doelf()
	if ctxt.HeadType == objabi.Hdarwin {
		ctxt.domacho()
	}
	ctxt.dostkcheck()
	if ctxt.HeadType == objabi.Hwindows {
		ctxt.dope()
		ctxt.windynrelocsyms()
	}
	if ctxt.HeadType == objabi.Haix {
		ctxt.doxcoff()
	}

	ctxt.addexport()
	thearch.Gentext(ctxt) // trampolines, call stubs, etc.
	ctxt.textbuildid()
	ctxt.textaddress()
	ctxt.pclntab()
	ctxt.findfunctab()
	ctxt.typelink()
	ctxt.symtab()
	ctxt.buildinfo()
	ctxt.dodata()
	order := ctxt.address()
	dwarfcompress(ctxt)
	filesize := ctxt.layout(order)

	// Write out the output file.
	// It is split into two parts (Asmb and Asmb2). The first
	// part writes most of the content (sections and segments),
	// for which we have computed the size and offset, in a
	// mmap'd region. The second part writes more content, for
	// which we don't know the size.
	var outputMmapped bool
	if ctxt.Arch.Family != sys.Wasm {
		// Don't mmap if we're building for Wasm. Wasm file
		// layout is very different so filesize is meaningless.
		err := ctxt.Out.Mmap(filesize)
		outputMmapped = err == nil
	}
	if outputMmapped {
		// Asmb will redirect symbols to the output file mmap, and relocations
		// will be applied directly there.
		thearch.Asmb(ctxt)
		ctxt.reloc()
		ctxt.Out.Munmap()
	} else {
		// If we don't mmap, we need to apply relocations before
		// writing out.
		ctxt.reloc()
		thearch.Asmb(ctxt)
	}
	thearch.Asmb2(ctxt)

	ctxt.undef()
	ctxt.hostlink()
	if ctxt.Debugvlog != 0 {
		ctxt.Logf("%d symbols\n", len(ctxt.Syms.Allsym))
		ctxt.Logf("%d liveness data\n", liveness)
	}
	ctxt.Bso.Flush()
	ctxt.archive()

	errorexit()
}

type Rpath struct {
	set bool
	val string
}

func (r *Rpath) Set(val string) error {
	r.set = true
	r.val = val
	return nil
}

func (r *Rpath) String() string {
	return r.val
}

func startProfile() {
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatalf("%v", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatalf("%v", err)
		}
		AtExit(pprof.StopCPUProfile)
	}
	if *memprofile != "" {
		if *memprofilerate != 0 {
			runtime.MemProfileRate = int(*memprofilerate)
		}
		f, err := os.Create(*memprofile)
		if err != nil {
			log.Fatalf("%v", err)
		}
		AtExit(func() {
			// Profile all outstanding allocations.
			runtime.GC()
			// compilebench parses the memory profile to extract memstats,
			// which are only written in the legacy pprof format.
			// See golang.org/issue/18641 and runtime/pprof/pprof.go:writeHeap.
			const writeLegacyFormat = 1
			if err := pprof.Lookup("heap").WriteTo(f, writeLegacyFormat); err != nil {
				log.Fatalf("%v", err)
			}
		})
	}
}
<commit_msg>[dev.link] cmd/link: default to new ELF host object loader<commit_after>// Inferno utils/6l/obj.c
// https://bitbucket.org/inferno-os/inferno-os/src/default/utils/6l/obj.c
//
//	Copyright © 1994-1999 Lucent Technologies Inc.  All rights reserved.
//	Portions Copyright © 1995-1997 C H Forsyth (forsyth@terzarima.net)
//	Portions Copyright © 1997-1999 Vita Nuova Limited
//	Portions Copyright © 2000-2007 Vita Nuova Holdings Limited (www.vitanuova.com)
//	Portions Copyright © 2004,2006 Bruce Ellis
//	Portions Copyright © 2005-2007 C H Forsyth (forsyth@terzarima.net)
//	Revisions Copyright © 2000-2007 Lucent Technologies Inc. and others
//	Portions Copyright © 2009 The Go Authors. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package ld

import (
	"bufio"
	"cmd/internal/objabi"
	"cmd/internal/sys"
	"cmd/link/internal/sym"
	"flag"
	"log"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
)

var (
	pkglistfornote []byte
	windowsgui     bool // writes a "GUI binary" instead of a "console binary"
	ownTmpDir      bool // set to true if tmp dir created by linker (e.g. no -tmpdir)
)

func init() {
	flag.Var(&rpath, "r", "set the ELF dynamic linker search `path` to dir1:dir2:...")
}

// Flags used by the linker. The exported flags are used by the architecture-specific packages.
var (
	flagBuildid = flag.String("buildid", "", "record `id` as Go toolchain build id")

	flagOutfile    = flag.String("o", "", "write output to `file`")
	flagPluginPath = flag.String("pluginpath", "", "full path name for plugin")

	flagInstallSuffix = flag.String("installsuffix", "", "set package directory `suffix`")
	flagDumpDep       = flag.Bool("dumpdep", false, "dump symbol dependency graph")
	flagRace          = flag.Bool("race", false, "enable race detector")
	flagMsan          = flag.Bool("msan", false, "enable MSan interface")

	flagFieldTrack = flag.String("k", "", "set field tracking `symbol`")
	flagLibGCC     = flag.String("libgcc", "", "compiler support lib for internal linking; use \"none\" to disable")
	flagTmpdir     = flag.String("tmpdir", "", "use `directory` for temporary files")

	flagExtld      = flag.String("extld", "", "use `linker` when linking in external mode")
	flagExtldflags = flag.String("extldflags", "", "pass `flags` to external linker")
	flagExtar      = flag.String("extar", "", "archive program for buildmode=c-archive")

	flagA           = flag.Bool("a", false, "disassemble output")
	FlagC           = flag.Bool("c", false, "dump call graph")
	FlagD           = flag.Bool("d", false, "disable dynamic executable")
	flagF           = flag.Bool("f", false, "ignore version mismatch")
	flagG           = flag.Bool("g", false, "disable go package data checks")
	flagH           = flag.Bool("h", false, "halt on error")
	flagN           = flag.Bool("n", false, "dump symbol table")
	FlagS           = flag.Bool("s", false, "disable symbol table")
	flagU           = flag.Bool("u", false, "reject unsafe packages")
	FlagW           = flag.Bool("w", false, "disable DWARF generation")
	Flag8           bool // use 64-bit addresses in symbol table
	flagInterpreter = flag.String("I", "", "use `linker` as ELF dynamic linker")
	FlagDebugTramp  = flag.Int("debugtramp", 0, "debug trampolines")
	FlagStrictDups  = flag.Int("strictdups", 0, "sanity check duplicate symbol contents during object file reading (1=warn 2=err).")
	FlagNewLdElf    = flag.Bool("newldelf", true, "ELF host obj load with new loader")

	FlagRound       = flag.Int("R", -1, "set address rounding `quantum`")
	FlagTextAddr    = flag.Int64("T", -1, "set text segment `address`")
	flagEntrySymbol = flag.String("E", "", "set `entry` symbol name")

	cpuprofile     = flag.String("cpuprofile", "", "write cpu profile to `file`")
	memprofile     = flag.String("memprofile", "", "write memory profile to `file`")
	memprofilerate = flag.Int64("memprofilerate", 0, "set runtime.MemProfileRate to `rate`")
)

// Main is the main entry point for the linker code.
func Main(arch *sys.Arch, theArch Arch) {
	thearch = theArch
	ctxt := linknew(arch)
	ctxt.Bso = bufio.NewWriter(os.Stdout)

	// For testing behavior of go command when tools crash silently.
	// Undocumented, not in standard flag parser to avoid
	// exposing in usage message.
	for _, arg := range os.Args {
		if arg == "-crash_for_testing" {
			os.Exit(2)
		}
	}

	final := gorootFinal()
	addstrdata1(ctxt, "runtime/internal/sys.DefaultGoroot="+final)
	addstrdata1(ctxt, "cmd/internal/objabi.defaultGOROOT="+final)

	// TODO(matloob): define these above and then check flag values here
	if ctxt.Arch.Family == sys.AMD64 && objabi.GOOS == "plan9" {
		flag.BoolVar(&Flag8, "8", false, "use 64-bit addresses in symbol table")
	}
	flagHeadType := flag.String("H", "", "set header `type`")
	flag.BoolVar(&ctxt.linkShared, "linkshared", false, "link against installed Go shared libraries")
	flag.Var(&ctxt.LinkMode, "linkmode", "set link `mode`")
	flag.Var(&ctxt.BuildMode, "buildmode", "set build `mode`")
	flag.BoolVar(&ctxt.compressDWARF, "compressdwarf", true, "compress DWARF if possible")
	objabi.Flagfn1("B", "add an ELF NT_GNU_BUILD_ID `note` when using ELF", addbuildinfo)
	objabi.Flagfn1("L", "add specified `directory` to library path", func(a string) { Lflag(ctxt, a) })
	objabi.AddVersionFlag() // -V
	objabi.Flagfn1("X", "add string value `definition` of the form importpath.name=value", func(s string) { addstrdata1(ctxt, s) })
	objabi.Flagcount("v", "print link trace", &ctxt.Debugvlog)
	objabi.Flagfn1("importcfg", "read import configuration from `file`", ctxt.readImportCfg)

	objabi.Flagparse(usage)

	switch *flagHeadType {
	case "":
	case "windowsgui":
		ctxt.HeadType = objabi.Hwindows
		windowsgui = true
	default:
		if err := ctxt.HeadType.Set(*flagHeadType); err != nil {
			Errorf(nil, "%v", err)
			usage()
		}
	}

	if objabi.Fieldtrack_enabled != 0 {
		ctxt.Reachparent = make(map[*sym.Symbol]*sym.Symbol)
	}
	checkStrictDups = *FlagStrictDups

	startProfile()
	if ctxt.BuildMode == BuildModeUnset {
		ctxt.BuildMode = BuildModeExe
	}

	if ctxt.BuildMode != BuildModeShared && flag.NArg() != 1 {
		usage()
	}

	if *flagOutfile == "" {
		*flagOutfile = "a.out"
		if ctxt.HeadType == objabi.Hwindows {
			*flagOutfile += ".exe"
		}
	}

	interpreter = *flagInterpreter

	libinit(ctxt) // creates outfile

	if ctxt.HeadType == objabi.Hunknown {
		ctxt.HeadType.Set(objabi.GOOS)
	}

	ctxt.computeTLSOffset()
	thearch.Archinit(ctxt)

	if ctxt.linkShared && !ctxt.IsELF {
		Exitf("-linkshared can only be used on elf systems")
	}

	if ctxt.Debugvlog != 0 {
		ctxt.Logf("HEADER = -H%d -T0x%x -R0x%x\n", ctxt.HeadType, uint64(*FlagTextAddr), uint32(*FlagRound))
	}

	switch ctxt.BuildMode {
	case BuildModeShared:
		for i := 0; i < flag.NArg(); i++ {
			arg := flag.Arg(i)
			parts := strings.SplitN(arg, "=", 2)
			var pkgpath, file string
			if len(parts) == 1 {
				pkgpath, file = "main", arg
			} else {
				pkgpath, file = parts[0], parts[1]
			}
			pkglistfornote = append(pkglistfornote, pkgpath...)
			pkglistfornote = append(pkglistfornote, '\n')
			addlibpath(ctxt, "command line", "command line", file, pkgpath, "")
		}
	case BuildModePlugin:
		addlibpath(ctxt, "command line", "command line", flag.Arg(0), *flagPluginPath, "")
	default:
		addlibpath(ctxt, "command line", "command line", flag.Arg(0), "main", "")
	}
	ctxt.loadlib()

	deadcode(ctxt)

	if *FlagNewLdElf {
		ctxt.linksetup()
	}

	ctxt.loadlibfull() // XXX do it here for now

	if !*FlagNewLdElf {
		ctxt.linksetupold()
	}
	ctxt.dostrdata()
	dwarfGenerateDebugInfo(ctxt)

	if objabi.Fieldtrack_enabled != 0 {
		fieldtrack(ctxt)
	}
	ctxt.mangleTypeSym()
	ctxt.callgraph()

	ctxt.doelf()
	if ctxt.HeadType == objabi.Hdarwin {
		ctxt.domacho()
	}
	ctxt.dostkcheck()
	if ctxt.HeadType == objabi.Hwindows {
		ctxt.dope()
		ctxt.windynrelocsyms()
	}
	if ctxt.HeadType == objabi.Haix {
		ctxt.doxcoff()
	}

	ctxt.addexport()
	thearch.Gentext(ctxt) // trampolines, call stubs, etc.
	ctxt.textbuildid()
	ctxt.textaddress()
	ctxt.pclntab()
	ctxt.findfunctab()
	ctxt.typelink()
	ctxt.symtab()
	ctxt.buildinfo()
	ctxt.dodata()
	order := ctxt.address()
	dwarfcompress(ctxt)
	filesize := ctxt.layout(order)

	// Write out the output file.
	// It is split into two parts (Asmb and Asmb2). The first
	// part writes most of the content (sections and segments),
	// for which we have computed the size and offset, in a
	// mmap'd region. The second part writes more content, for
	// which we don't know the size.
	var outputMmapped bool
	if ctxt.Arch.Family != sys.Wasm {
		// Don't mmap if we're building for Wasm. Wasm file
		// layout is very different so filesize is meaningless.
		err := ctxt.Out.Mmap(filesize)
		outputMmapped = err == nil
	}
	if outputMmapped {
		// Asmb will redirect symbols to the output file mmap, and relocations
		// will be applied directly there.
		thearch.Asmb(ctxt)
		ctxt.reloc()
		ctxt.Out.Munmap()
	} else {
		// If we don't mmap, we need to apply relocations before
		// writing out.
		ctxt.reloc()
		thearch.Asmb(ctxt)
	}
	thearch.Asmb2(ctxt)

	ctxt.undef()
	ctxt.hostlink()
	if ctxt.Debugvlog != 0 {
		ctxt.Logf("%d symbols\n", len(ctxt.Syms.Allsym))
		ctxt.Logf("%d liveness data\n", liveness)
	}
	ctxt.Bso.Flush()
	ctxt.archive()

	errorexit()
}

type Rpath struct {
	set bool
	val string
}

func (r *Rpath) Set(val string) error {
	r.set = true
	r.val = val
	return nil
}

func (r *Rpath) String() string {
	return r.val
}

func startProfile() {
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatalf("%v", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatalf("%v", err)
		}
		AtExit(pprof.StopCPUProfile)
	}
	if *memprofile != "" {
		if *memprofilerate != 0 {
			runtime.MemProfileRate = int(*memprofilerate)
		}
		f, err := os.Create(*memprofile)
		if err != nil {
			log.Fatalf("%v", err)
		}
		AtExit(func() {
			// Profile all outstanding allocations.
			runtime.GC()
			// compilebench parses the memory profile to extract memstats,
			// which are only written in the legacy pprof format.
			// See golang.org/issue/18641 and runtime/pprof/pprof.go:writeHeap.
			const writeLegacyFormat = 1
			if err := pprof.Lookup("heap").WriteTo(f, writeLegacyFormat); err != nil {
				log.Fatalf("%v", err)
			}
		})
	}
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package comm

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/hyperledger/fabric/gossip/common"
	"github.com/hyperledger/fabric/gossip/metrics"
	"github.com/hyperledger/fabric/gossip/protoext"
	"github.com/hyperledger/fabric/gossip/util"
	proto "github.com/hyperledger/fabric/protos/gossip"
	"github.com/pkg/errors"
	"go.uber.org/zap/zapcore"
	"google.golang.org/grpc"
)

type handler func(message *protoext.SignedGossipMessage)

type blockingBehavior bool

const (
	blockingSend    = blockingBehavior(true)
	nonBlockingSend = blockingBehavior(false)
)

type connFactory interface {
	createConnection(endpoint string, pkiID common.PKIidType) (*connection, error)
}

type connectionStore struct {
	config           ConnConfig
	logger           util.Logger            // logger
	isClosing        bool                   // whether this connection store is shutting down
	connFactory      connFactory            // creates a connection to remote peer
	sync.RWMutex                            // synchronize access to shared variables
	pki2Conn         map[string]*connection // mapping between pkiID to connections
	destinationLocks map[string]*sync.Mutex //mapping between pkiIDs and locks,
	// used to prevent concurrent connection establishment to the same remote endpoint
}

func newConnStore(connFactory connFactory, logger util.Logger, config ConnConfig) *connectionStore {
	return &connectionStore{
		connFactory:      connFactory,
		isClosing:        false,
		pki2Conn:         make(map[string]*connection),
		destinationLocks: make(map[string]*sync.Mutex),
		logger:           logger,
		config:           config,
	}
}

func (cs *connectionStore) getConnection(peer *RemotePeer) (*connection, error) {
	cs.RLock()
	isClosing := cs.isClosing
	cs.RUnlock()

	if isClosing {
		return nil, fmt.Errorf("Shutting down")
	}

	pkiID := peer.PKIID
	endpoint := peer.Endpoint

	cs.Lock()
	destinationLock, hasConnected := cs.destinationLocks[string(pkiID)]
	if !hasConnected {
		destinationLock = &sync.Mutex{}
		cs.destinationLocks[string(pkiID)] = destinationLock
	}
	cs.Unlock()

	destinationLock.Lock()

	cs.RLock()
	conn, exists := cs.pki2Conn[string(pkiID)]
	if exists {
		cs.RUnlock()
		destinationLock.Unlock()
		return conn, nil
	}
	cs.RUnlock()

	createdConnection, err := cs.connFactory.createConnection(endpoint, pkiID)

	destinationLock.Unlock()

	cs.RLock()
	isClosing = cs.isClosing
	cs.RUnlock()
	if isClosing {
		return nil, fmt.Errorf("ConnStore is closing")
	}

	cs.Lock()
	delete(cs.destinationLocks, string(pkiID))
	defer cs.Unlock()

	// check again, maybe someone connected to us during the connection creation?
	conn, exists = cs.pki2Conn[string(pkiID)]

	if exists {
		if createdConnection != nil {
			createdConnection.close()
		}
		return conn, nil
	}

	// no one connected to us AND we failed connecting!
	if err != nil {
		return nil, err
	}

	// at this point in the code, we created a connection to a remote peer
	conn = createdConnection
	cs.pki2Conn[string(createdConnection.pkiID)] = conn

	go conn.serviceConnection()

	return conn, nil
}

func (cs *connectionStore) connNum() int {
	cs.RLock()
	defer cs.RUnlock()
	return len(cs.pki2Conn)
}

func (cs *connectionStore) closeConn(peer *RemotePeer) {
	cs.Lock()
	defer cs.Unlock()

	if conn, exists := cs.pki2Conn[string(peer.PKIID)]; exists {
		conn.close()
		delete(cs.pki2Conn, string(conn.pkiID))
	}
}

func (cs *connectionStore) shutdown() {
	cs.Lock()
	cs.isClosing = true
	pkiIds2conn := cs.pki2Conn

	var connections2Close []*connection
	for _, conn := range pkiIds2conn {
		connections2Close = append(connections2Close, conn)
	}
	cs.Unlock()

	wg := sync.WaitGroup{}
	for _, conn := range connections2Close {
		wg.Add(1)
		go func(conn *connection) {
			cs.closeByPKIid(conn.pkiID)
			wg.Done()
		}(conn)
	}
	wg.Wait()
}

func (cs *connectionStore) onConnected(serverStream proto.Gossip_GossipStreamServer,
	connInfo *protoext.ConnectionInfo, metrics *metrics.CommMetrics) *connection {
	cs.Lock()
	defer cs.Unlock()

	if c, exists := cs.pki2Conn[string(connInfo.ID)]; exists {
		c.close()
	}

	return cs.registerConn(connInfo, serverStream, metrics)
}

func (cs *connectionStore) registerConn(connInfo *protoext.ConnectionInfo,
	serverStream proto.Gossip_GossipStreamServer, metrics *metrics.CommMetrics) *connection {
	conn := newConnection(nil, nil, nil, serverStream, metrics, cs.config)
	conn.pkiID = connInfo.ID
	conn.info = connInfo
	conn.logger = cs.logger
	cs.pki2Conn[string(connInfo.ID)] = conn
	return conn
}

func (cs *connectionStore) closeByPKIid(pkiID common.PKIidType) {
	cs.Lock()
	defer cs.Unlock()
	if conn, exists := cs.pki2Conn[string(pkiID)]; exists {
		conn.close()
		delete(cs.pki2Conn, string(pkiID))
	}
}

func newConnection(cl proto.GossipClient, c *grpc.ClientConn, cs proto.Gossip_GossipStreamClient,
	ss proto.Gossip_GossipStreamServer, metrics *metrics.CommMetrics, config ConnConfig) *connection {
	connection := &connection{
		metrics:      metrics,
		outBuff:      make(chan *msgSending, config.SendBuffSize),
		cl:           cl,
		conn:         c,
		clientStream: cs,
		serverStream: ss,
		stopFlag:     int32(0),
		stopChan:     make(chan struct{}, 1),
		recvBuffSize: config.RecvBuffSize,
	}
	return connection
}

// ConnConfig is the configuration required to initialize a new conn
type ConnConfig struct {
	RecvBuffSize int
	SendBuffSize int
}

type connection struct {
	recvBuffSize int
	metrics      *metrics.CommMetrics
	cancel       context.CancelFunc
	info         *protoext.ConnectionInfo
	outBuff      chan *msgSending
	logger       util.Logger                     // logger
	pkiID        common.PKIidType                // pkiID of the remote endpoint
	handler      handler                         // function to invoke upon a message reception
	conn         *grpc.ClientConn                // gRPC connection to remote endpoint
	cl           proto.GossipClient              // gRPC stub of remote endpoint
	clientStream proto.Gossip_GossipStreamClient // client-side stream to remote endpoint
	serverStream proto.Gossip_GossipStreamServer // server-side stream to remote endpoint
	stopFlag     int32                           // indicates whether this connection is in process of stopping
	stopChan     chan struct{}                   // a method to stop the server-side gRPC call from a different go-routine
	sync.RWMutex                                 // synchronizes access to shared variables
	stopWG       sync.WaitGroup                  // a method to wait for stream activity to stop before closing it
}

func (conn *connection) close() {
	if conn.toDie() {
		return
	}

	amIFirst := atomic.CompareAndSwapInt32(&conn.stopFlag, int32(0), int32(1))
	if !amIFirst {
		return
	}

	conn.stopChan <- struct{}{}

	conn.drainOutputBuffer()
	conn.Lock()
	defer conn.Unlock()

	if conn.clientStream != nil {
		conn.stopWG.Wait()
		conn.clientStream.CloseSend()
	}
	if conn.conn != nil {
		conn.conn.Close()
	}

	if conn.cancel != nil {
		conn.cancel()
	}
}

func (conn *connection) toDie() bool {
	return atomic.LoadInt32(&(conn.stopFlag)) == int32(1)
}

func (conn *connection) send(msg *protoext.SignedGossipMessage, onErr func(error), shouldBlock blockingBehavior) {
	if conn.toDie() {
		conn.logger.Debug("Aborting send() to ", conn.info.Endpoint, "because connection is closing")
		return
	}

	m := &msgSending{
		envelope: msg.Envelope,
		onErr:    onErr,
	}

	if len(conn.outBuff) == cap(conn.outBuff) {
		if conn.logger.IsEnabledFor(zapcore.DebugLevel) {
			conn.logger.Debug("Buffer to", conn.info.Endpoint, "overflowed, dropping message", msg.String())
			conn.metrics.BufferOverflow.Add(1)
		}
		if !shouldBlock {
			return
		}
	}

	conn.outBuff <- m
}

func (conn *connection) serviceConnection() error {
	errChan := make(chan error, 1)
	msgChan := make(chan *protoext.SignedGossipMessage, conn.recvBuffSize)
	quit := make(chan struct{})
	// Call stream.Recv() asynchronously in readFromStream(),
	// and wait for either the Recv() call to end,
	// or a signal to close the connection, which exits
	// the method and makes the Recv() call to fail in the
	// readFromStream() method
	go conn.readFromStream(errChan, quit, msgChan)

	conn.stopWG.Add(1) // wait for write to finish before closing it
	go conn.writeToStream()

	for !conn.toDie() {
		select {
		case stop := <-conn.stopChan:
			conn.logger.Debug("Closing reading from stream")
			conn.stopChan <- stop
			return nil
		case err := <-errChan:
			return err
		case msg := <-msgChan:
			conn.handler(msg)
		}
	}
	return nil
}

func (conn *connection) writeToStream() {
	defer conn.stopWG.Done()
	for !conn.toDie() {
		stream := conn.getStream()
		if stream == nil {
			conn.logger.Error(conn.pkiID, "Stream is nil, aborting!")
			return
		}
		select {
		case m := <-conn.outBuff:
			err := stream.Send(m.envelope)
			if err != nil {
				go m.onErr(err)
				return
			}
			conn.metrics.SentMessages.Add(1)
		case stop := <-conn.stopChan:
			conn.logger.Debug("Closing writing to stream")
			conn.stopChan <- stop
			return
		}
	}
}

func (conn *connection) drainOutputBuffer() {
	// Drain the output buffer
	for len(conn.outBuff) > 0 {
		<-conn.outBuff
	}
}

func (conn *connection) readFromStream(errChan chan error, quit chan struct{}, msgChan chan *protoext.SignedGossipMessage) {
	for !conn.toDie() {
		stream := conn.getStream()
		if stream == nil {
			conn.logger.Error(conn.pkiID, "Stream is nil, aborting!")
			errChan <- fmt.Errorf("Stream is nil")
			return
		}
		envelope, err := stream.Recv()
		if conn.toDie() {
			conn.logger.Debug(conn.pkiID, "canceling read because closing")
			return
		}
		if err != nil {
			errChan <- err
			conn.logger.Debugf("Got error, aborting: %v", err)
			return
		}
		conn.metrics.ReceivedMessages.Add(1)
		msg, err := protoext.EnvelopeToGossipMessage(envelope)
		if err != nil {
			errChan <- err
			conn.logger.Warningf("Got error, aborting: %v", err)
		}
		select {
		case msgChan <- msg:
		case <-quit:
			return
		}
	}
}

func (conn *connection) getStream() stream {
	conn.Lock()
	defer conn.Unlock()

	if conn.clientStream != nil && conn.serverStream != nil {
		e := errors.New("Both client and server stream are not nil, something went wrong")
		conn.logger.Errorf("%+v", e)
	}

	if conn.clientStream != nil {
		return conn.clientStream
	}

	if conn.serverStream != nil {
		return conn.serverStream
	}

	return nil
}

type msgSending struct {
	envelope *proto.Envelope
	onErr    func(error)
}
<commit_msg>FAB-15032 avoid race between len() and channel read<commit_after>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package comm

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/hyperledger/fabric/gossip/common"
	"github.com/hyperledger/fabric/gossip/metrics"
	"github.com/hyperledger/fabric/gossip/protoext"
	"github.com/hyperledger/fabric/gossip/util"
	proto "github.com/hyperledger/fabric/protos/gossip"
	"github.com/pkg/errors"
	"google.golang.org/grpc"
)

type handler func(message *protoext.SignedGossipMessage)

type blockingBehavior bool

const (
	blockingSend    = blockingBehavior(true)
	nonBlockingSend = blockingBehavior(false)
)

type connFactory interface {
	createConnection(endpoint string, pkiID common.PKIidType) (*connection, error)
}

type connectionStore struct {
	config           ConnConfig
	logger           util.Logger            // logger
	isClosing        bool                   // whether this connection store is shutting down
	connFactory      connFactory            // creates a connection to remote peer
	sync.RWMutex                            // synchronize access to shared variables
	pki2Conn         map[string]*connection // mapping between pkiID to connections
	destinationLocks map[string]*sync.Mutex //mapping between pkiIDs and locks,
	// used to prevent concurrent connection establishment to the same remote endpoint
}

func newConnStore(connFactory connFactory, logger util.Logger, config ConnConfig) *connectionStore {
	return &connectionStore{
		connFactory:      connFactory,
		isClosing:        false,
		pki2Conn:         make(map[string]*connection),
		destinationLocks: make(map[string]*sync.Mutex),
		logger:           logger,
		config:           config,
	}
}

func (cs *connectionStore) getConnection(peer *RemotePeer) (*connection, error) {
	cs.RLock()
	isClosing := cs.isClosing
	cs.RUnlock()

	if isClosing {
		return nil, fmt.Errorf("Shutting down")
	}

	pkiID := peer.PKIID
	endpoint := peer.Endpoint

	cs.Lock()
	destinationLock, hasConnected := cs.destinationLocks[string(pkiID)]
	if !hasConnected {
		destinationLock = &sync.Mutex{}
		cs.destinationLocks[string(pkiID)] = destinationLock
	}
	cs.Unlock()

	destinationLock.Lock()

	cs.RLock()
	conn, exists := cs.pki2Conn[string(pkiID)]
	if exists {
		cs.RUnlock()
		destinationLock.Unlock()
		return conn, nil
	}
	cs.RUnlock()

	createdConnection, err := cs.connFactory.createConnection(endpoint, pkiID)

	destinationLock.Unlock()

	cs.RLock()
	isClosing = cs.isClosing
	cs.RUnlock()
	if isClosing {
		return nil, fmt.Errorf("ConnStore is closing")
	}

	cs.Lock()
	delete(cs.destinationLocks, string(pkiID))
	defer cs.Unlock()

	// check again, maybe someone connected to us during the connection creation?
	conn, exists = cs.pki2Conn[string(pkiID)]

	if exists {
		if createdConnection != nil {
			createdConnection.close()
		}
		return conn, nil
	}

	// no one connected to us AND we failed connecting!
	if err != nil {
		return nil, err
	}

	// at this point in the code, we created a connection to a remote peer
	conn = createdConnection
	cs.pki2Conn[string(createdConnection.pkiID)] = conn

	go conn.serviceConnection()

	return conn, nil
}

func (cs *connectionStore) connNum() int {
	cs.RLock()
	defer cs.RUnlock()
	return len(cs.pki2Conn)
}

func (cs *connectionStore) closeConn(peer *RemotePeer) {
	cs.Lock()
	defer cs.Unlock()

	if conn, exists := cs.pki2Conn[string(peer.PKIID)]; exists {
		conn.close()
		delete(cs.pki2Conn, string(conn.pkiID))
	}
}

func (cs *connectionStore) shutdown() {
	cs.Lock()
	cs.isClosing = true
	pkiIds2conn := cs.pki2Conn

	var connections2Close []*connection
	for _, conn := range pkiIds2conn {
		connections2Close = append(connections2Close, conn)
	}
	cs.Unlock()

	wg := sync.WaitGroup{}
	for _, conn := range connections2Close {
		wg.Add(1)
		go func(conn *connection) {
			cs.closeByPKIid(conn.pkiID)
			wg.Done()
		}(conn)
	}
	wg.Wait()
}

func (cs *connectionStore) onConnected(serverStream proto.Gossip_GossipStreamServer,
	connInfo *protoext.ConnectionInfo, metrics *metrics.CommMetrics) *connection {
	cs.Lock()
	defer cs.Unlock()

	if c, exists := cs.pki2Conn[string(connInfo.ID)]; exists {
		c.close()
	}

	return cs.registerConn(connInfo, serverStream, metrics)
}

func (cs *connectionStore) registerConn(connInfo *protoext.ConnectionInfo,
	serverStream proto.Gossip_GossipStreamServer, metrics *metrics.CommMetrics) *connection {
	conn := newConnection(nil, nil, nil, serverStream, metrics, cs.config)
	conn.pkiID = connInfo.ID
	conn.info = connInfo
	conn.logger = cs.logger
	cs.pki2Conn[string(connInfo.ID)] = conn
	return conn
}

func (cs *connectionStore) closeByPKIid(pkiID common.PKIidType) {
	cs.Lock()
	defer cs.Unlock()
	if conn, exists := cs.pki2Conn[string(pkiID)]; exists {
		conn.close()
		delete(cs.pki2Conn, string(pkiID))
	}
}

func newConnection(cl proto.GossipClient, c *grpc.ClientConn, cs proto.Gossip_GossipStreamClient,
	ss proto.Gossip_GossipStreamServer, metrics *metrics.CommMetrics, config ConnConfig) *connection {
	connection := &connection{
		metrics:      metrics,
		outBuff:      make(chan *msgSending, config.SendBuffSize),
		cl:           cl,
		conn:         c,
		clientStream: cs,
		serverStream: ss,
		stopFlag:     int32(0),
		stopChan:     make(chan struct{}, 1),
		recvBuffSize: config.RecvBuffSize,
	}
	return connection
}

// ConnConfig is the configuration required to initialize a new conn
type ConnConfig struct {
	RecvBuffSize int
	SendBuffSize int
}

type connection struct {
	recvBuffSize int
	metrics      *metrics.CommMetrics
	cancel       context.CancelFunc
	info         *protoext.ConnectionInfo
	outBuff      chan *msgSending
	logger       util.Logger                     // logger
	pkiID        common.PKIidType                // pkiID of the remote endpoint
	handler      handler                         // function to invoke upon a message reception
	conn         *grpc.ClientConn                // gRPC connection to remote endpoint
	cl           proto.GossipClient              // gRPC stub of remote endpoint
	clientStream proto.Gossip_GossipStreamClient // client-side stream to remote endpoint
	serverStream proto.Gossip_GossipStreamServer // server-side stream to remote endpoint
	stopFlag     int32                           // indicates whether this connection is in process of stopping
	stopChan     chan struct{}                   // a method to stop the server-side gRPC call from a different go-routine
	sync.RWMutex                                 // synchronizes access to shared variables
	stopWG       sync.WaitGroup                  // a method to wait for stream activity to stop before closing it
}

func (conn *connection) close() {
	if conn.toDie() {
		return
	}

	amIFirst := atomic.CompareAndSwapInt32(&conn.stopFlag, int32(0), int32(1))
	if !amIFirst {
		return
	}

	conn.stopChan <- struct{}{}

	conn.drainOutputBuffer()
	conn.Lock()
	defer conn.Unlock()

	if conn.clientStream != nil {
		conn.stopWG.Wait()
		conn.clientStream.CloseSend()
	}
	if conn.conn != nil {
		conn.conn.Close()
	}

	if conn.cancel != nil {
		conn.cancel()
	}
}

func (conn *connection) toDie() bool {
	return atomic.LoadInt32(&(conn.stopFlag)) == int32(1)
}

func (conn *connection) send(msg *protoext.SignedGossipMessage, onErr func(error), shouldBlock blockingBehavior) {
	if conn.toDie() {
		conn.logger.Debugf("Aborting send() to %s because connection is closing", conn.info.Endpoint)
		return
	}

	m := &msgSending{
		envelope: msg.Envelope,
		onErr:    onErr,
	}

	select {
	case conn.outBuff <- m:
		// room in channel, successfully sent message, nothing to do
	default: // did not send
		if shouldBlock {
			conn.outBuff <- m // try again, and wait to send
		} else {
			conn.metrics.BufferOverflow.Add(1)
			conn.logger.Debugf("Buffer to %s overflowed, dropping message %s", conn.info.Endpoint, msg)
		}
	}
}

func (conn *connection) serviceConnection() error {
	errChan := make(chan error, 1)
	msgChan := make(chan *protoext.SignedGossipMessage, conn.recvBuffSize)
	quit := make(chan struct{})
	// Call stream.Recv() asynchronously in readFromStream(),
	// and wait for either the Recv() call to end,
	// or a signal to close the connection, which exits
	// the method and makes the Recv() call to fail in the
	// readFromStream() method
	go conn.readFromStream(errChan, quit, msgChan)

	conn.stopWG.Add(1) // wait for write to finish before closing it
	go conn.writeToStream()

	for !conn.toDie() {
		select {
		case stop := <-conn.stopChan:
			conn.logger.Debug("Closing reading from stream")
			conn.stopChan <- stop
			return nil
		case err := <-errChan:
			return err
		case msg := <-msgChan:
			conn.handler(msg)
		}
	}
	return nil
}

func (conn *connection) writeToStream() {
	defer conn.stopWG.Done()
	for !conn.toDie() {
		stream := conn.getStream()
		if stream == nil {
			conn.logger.Error(conn.pkiID, "Stream is nil, aborting!")
			return
		}
		select {
		case m := <-conn.outBuff:
			err := stream.Send(m.envelope)
			if err != nil {
				go m.onErr(err)
				return
			}
			conn.metrics.SentMessages.Add(1)
		case stop := <-conn.stopChan:
			conn.logger.Debug("Closing writing to stream")
			conn.stopChan <- stop
			return
		}
	}
}

func (conn *connection) drainOutputBuffer() {
	// Read from the buffer until it is empty.
	// There may be multiple concurrent readers.
	for {
		select {
		case <-conn.outBuff:
		default:
			return
		}
	}
}

func (conn *connection) readFromStream(errChan chan error, quit chan struct{}, msgChan chan *protoext.SignedGossipMessage) {
	for !conn.toDie() {
		stream := conn.getStream()
		if stream == nil {
			conn.logger.Error(conn.pkiID, "Stream is nil, aborting!")
			errChan <- fmt.Errorf("Stream is nil")
			return
		}
		envelope, err := stream.Recv()
		if conn.toDie() {
			conn.logger.Debug(conn.pkiID, "canceling read because closing")
			return
		}
		if err != nil {
			errChan <- err
			conn.logger.Debugf("Got error, aborting: %v", err)
			return
		}
		conn.metrics.ReceivedMessages.Add(1)
		msg, err := protoext.EnvelopeToGossipMessage(envelope)
		if err != nil {
			errChan <- err
			conn.logger.Warningf("Got error, aborting: %v", err)
		}
		select {
		case msgChan <- msg:
		case <-quit:
			return
		}
	}
}

func (conn *connection) getStream() stream {
	conn.Lock()
	defer conn.Unlock()

	if conn.clientStream != nil && conn.serverStream != nil {
		e := errors.New("Both client and server stream are not nil, something went wrong")
		conn.logger.Errorf("%+v", e)
	}

	if conn.clientStream != nil {
		return conn.clientStream
	}

	if conn.serverStream != nil {
		return conn.serverStream
	}

	return nil
}

type msgSending struct {
	envelope *proto.Envelope
	onErr    func(error)
}
<|endoftext|>


<commit_before>package main

const primeIndexSize = 1 << 24

type Octree struct {
	N        int
	leafSize int
	mask     int
	p        [][]uint16
	v        []uint16
}

// N must be power of 2.
func NewOctree(N int) *Octree {
	indexSize := N
	if N > primeIndexSize {
		indexSize = primeIndexSize
	}
	return &Octree{
		N:        N,
		leafSize: N / indexSize,
		mask:     indexSize - 1,
		p:        make([][]uint16, indexSize),
		v:        make([]uint16, indexSize),
	}
}

func step(p, base [3]int, l int, index int64) (nbase [3]int, nl int, nindex int64) {
	nbase, nl, nindex = base, l>>1, index
	for i := 0; i < 3; i++ {
		nindex <<= 1
		if p[i] >= nbase[i]+nl {
			nindex++
			nbase[i] += nl
		}
	}
	return
}

func (t *Octree) GetV(x, y, z int) uint16 {
	p := [3]int{x, y, z}
	base := [3]int{0, 0, 0}
	var index int64
	l := t.N
	count := uint(1)
	base, l, index = step(p, base, l, index)
	if t.p != nil {
		panic("extended part of octree is not implemented")
	}
	cur := t.v[int(index&int64(t.mask))]
	if cur == 0 {
		return 0
	}

	// We need to check the neighbour if it has the value of 0 or not.
	// If no, we need to go deeper.
	// This is because we reuse the cell after an expansion.
	another := index + (1 << (count * 3))
	if t.v[int(another&int64(t.mask))] == 0 {
		// This is the leaf
		return cur
	}

	// The node is not leaf
	panic("not implemented")
}
<commit_msg>Check index of bounds<commit_after>package main

const primeIndexSize = 1 << 24

type Octree struct {
	N        int
	leafSize int
	mask     int64
	p        [][]uint16
	v        []uint16
}

// N must be power of 2.
func NewOctree(N int) *Octree {
	indexSize := N
	if N > primeIndexSize {
		indexSize = primeIndexSize
	}
	return &Octree{
		N:        N,
		leafSize: N / indexSize,
		mask:     int64(indexSize - 1),
		p:        make([][]uint16, indexSize),
		v:        make([]uint16, indexSize),
	}
}

func step(p, base [3]int, l int, index int64) (nbase [3]int, nl int, nindex int64) {
	nbase, nl, nindex = base, l>>1, index
	for i := 0; i < 3; i++ {
		nindex <<= 1
		if p[i] >= nbase[i]+nl {
			nindex++
			nbase[i] += nl
		}
	}
	return
}

func (t *Octree) GetV(x, y, z int) uint16 {
	p := [3]int{x, y, z}
	base := [3]int{0, 0, 0}
	var index int64
	l := t.N
	count := uint(1)
	base, l, index = step(p, base, l, index)
	if t.p != nil {
		panic("extended part of octree is not implemented")
	}
	cur := t.v[int(index&int64(t.mask))]
	if cur == 0 {
		return 0
	}

	// We need to check the neighbour if it has the value of 0 or not.
	// If no, we need to go deeper.
	// This is because we reuse the cell after an expansion.
	another := index + (1 << (count * 3))
	if another > t.mask || t.v[int(another&t.mask)] == 0 {
		// This is the leaf
		return cur
	}

	// The node is not leaf
	panic("not implemented")
}
<|endoftext|>