<commit_before>package trafcacc

import (
	"bytes"
	"encoding/gob"
	"errors"
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Sirupsen/logrus"
)

type upstream struct {
	uuid  uint64
	proto string
	alive int64
	grp   int

	// status recorder
	sent uint64
	recv uint64

	// tcp only
	encoder *gob.Encoder
	decoder *gob.Decoder

	// udp only (server)
	udpconn *net.UDPConn
	udpaddr *net.UDPAddr

	// dialer only
	conn net.Conn
	addr string
}

func (u *upstream) send(cmd cmd) error {
	p := &packet{Cmd: cmd}
	return u.sendpacket(p)
}

func (u *upstream) sendpacket(p *packet) error {
	atomic.AddUint64(&u.sent, uint64(len(p.Buf)))
	switch u.proto {
	case tcp:
		err := u.encoder.Encode(p)
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
				"cmd":   p.Cmd,
				"proto": u.proto,
			}).Warnln("send upstream cmd error")
		}
		return err
	case udp:
		var buf bytes.Buffer
		if err := gob.NewEncoder(&buf).Encode(p); err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
				"cmd":   p.Cmd,
				"proto": u.proto,
			}).Warnln("send upstream cmd error")
		}
		var err error
		if u.udpaddr != nil { // server
			_, err = u.udpconn.WriteToUDP(buf.Bytes(), u.udpaddr)
		} else if u.conn != nil { // dialer
			_, err = u.conn.Write(buf.Bytes())
		} else {
			logrus.WithFields(logrus.Fields{
				"upstream": u,
			}).Warnln("upstream is not there")
		}
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
				"cmd":   p.Cmd,
				"proto": u,
			}).Warnln("send upstream error")
		}
		return err
	}
	return errors.New("send to unknown upstream protocol")
}

func (u *upstream) close() {
	if u.conn != nil {
		u.conn.Close()
		u.conn = nil
	}
	if u.udpconn != nil {
		u.udpconn.Close()
		u.udpconn = nil
	}
}

func (u *upstream) isAlive() bool {
	return keepalive > time.Duration(time.Now().UnixNano()-atomic.LoadInt64(&u.alive))
}

type streampool struct {
	*sync.Cond
	mux      *sync.RWMutex
	pool     []*upstream
	atomicid uint64
	alive    bool

	// for pick up streams
	tcpool                 []*upstream
	udpool                 []*upstream
	alived                 []*upstream
	tcplen, udplen, alvlen int
}

func newStreamPool() *streampool {
	mux := &sync.RWMutex{}
	pl := &streampool{
		// use RWMutex
		Cond: sync.NewCond(mux),
		mux:  mux,
	}

	go pl.updateloop()
	return pl
}

func (pool *streampool) append(u *upstream, grp int) {
	pool.L.Lock()
	defer func() {
		pool.L.Unlock()
		pool.Broadcast()
	}()

	if u.uuid != 0 {
		for _, v := range pool.pool {
			if v.uuid == u.uuid {
				return
			}
		}
	}
	u.uuid = atomic.AddUint64(&pool.atomicid, 1)
	u.grp = grp
	pool.pool = append(pool.pool, u)

	pool.tcpool = pool.ensureCap(pool.tcpool)
	pool.udpool = pool.ensureCap(pool.udpool)
	pool.alived = pool.ensureCap(pool.alived)
}

func (pool *streampool) ensureCap(pl []*upstream) []*upstream {
	if cap(pl) < len(pool.pool) {
		ups := make([]*upstream, cap(pl)+10)
		for k, v := range pl {
			ups[k] = v
		}
		return ups
	}
	return pl
}

func (pool *streampool) pickupstreams() []*upstream {
	pool.waitforalive()

	// pick udp and tcp equally
	pool.mux.RLock()
	defer pool.mux.RUnlock()

	// pick one of each

	rn := int(rand.Int31())

	switch {
	case pool.tcplen > 0 && pool.udplen > 0:
		// pick one of each
		return []*upstream{
			pool.udpool[rn%pool.udplen],
			pool.tcpool[rn%pool.tcplen],
			//pool.udpool[(rn+1)%pool.udplen],
		}
	case pool.tcplen == 0 || pool.udplen == 0:
		// pick 1-2 alived
		return []*upstream{
			pool.alived[rn%pool.alvlen],
			pool.alived[(rn+1)%pool.alvlen],
		}
	}
	logrus.Warnln("no upstream avalible for pick")
	return nil

}

func (pool *streampool) waitforalive() {
	pool.L.Lock()
	for !pool.alive {
		pool.Wait()
	}
	pool.L.Unlock()
}

func (pool *streampool) updateloop() {
	for {
		pool.L.Lock()
		for !pool.updatealive() {
			pool.Wait()
		}
		pool.L.Unlock()
		pool.Broadcast()<commit_msg>roll back to pick 2 upstreams<commit_after>package trafcacc

import (
	"bytes"
	"encoding/gob"
	"errors"
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Sirupsen/logrus"
)

type upstream struct {
	uuid  uint64
	proto string
	alive int64
	grp   int

	// status recorder
	sent uint64
	recv uint64

	// tcp only
	encoder *gob.Encoder
	decoder *gob.Decoder

	// udp only (server)
	udpconn *net.UDPConn
	udpaddr *net.UDPAddr

	// dialer only
	conn net.Conn
	addr string
}

func (u *upstream) send(cmd cmd) error {
	p := &packet{Cmd: cmd}
	return u.sendpacket(p)
}

func (u *upstream) sendpacket(p *packet) error {
	atomic.AddUint64(&u.sent, uint64(len(p.Buf)))
	switch u.proto {
	case tcp:
		err := u.encoder.Encode(p)
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
				"cmd":   p.Cmd,
				"proto": u.proto,
			}).Warnln("send upstream cmd error")
		}
		return err
	case udp:
		var buf bytes.Buffer
		if err := gob.NewEncoder(&buf).Encode(p); err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
				"cmd":   p.Cmd,
				"proto": u.proto,
			}).Warnln("send upstream cmd error")
		}
		var err error
		if u.udpaddr != nil { // server
			_, err = u.udpconn.WriteToUDP(buf.Bytes(), u.udpaddr)
		} else if u.conn != nil { // dialer
			_, err = u.conn.Write(buf.Bytes())
		} else {
			logrus.WithFields(logrus.Fields{
				"upstream": u,
			}).Warnln("upstream is not there")
		}
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"error": err,
				"cmd":   p.Cmd,
				"proto": u,
			}).Warnln("send upstream error")
		}
		return err
	}
	return errors.New("send to unknown upstream protocol")
}

func (u *upstream) close() {
	if u.conn != nil {
		u.conn.Close()
		u.conn = nil
	}
	if u.udpconn != nil {
		u.udpconn.Close()
		u.udpconn = nil
	}
}

func (u *upstream) isAlive() bool {
	return keepalive > time.Duration(time.Now().UnixNano()-atomic.LoadInt64(&u.alive))
}

type streampool struct {
	*sync.Cond
	mux      *sync.RWMutex
	pool     []*upstream
	atomicid uint64
	alive    bool

	// for pick up streams
	tcpool                 []*upstream
	udpool                 []*upstream
	alived                 []*upstream
	tcplen, udplen, alvlen int
}

func newStreamPool() *streampool {
	mux := &sync.RWMutex{}
	pl := &streampool{
		// use RWMutex
		Cond: sync.NewCond(mux),
		mux:  mux,
	}

	go pl.updateloop()
	return pl
}

func (pool *streampool) append(u *upstream, grp int) {
	pool.L.Lock()
	defer func() {
		pool.L.Unlock()
		pool.Broadcast()
	}()

	if u.uuid != 0 {
		for _, v := range pool.pool {
			if v.uuid == u.uuid {
				return
			}
		}
	}
	u.uuid = atomic.AddUint64(&pool.atomicid, 1)
	u.grp = grp
	pool.pool = append(pool.pool, u)

	pool.tcpool = pool.ensureCap(pool.tcpool)
	pool.udpool = pool.ensureCap(pool.udpool)
	pool.alived = pool.ensureCap(pool.alived)
}

func (pool *streampool) ensureCap(pl []*upstream) []*upstream {
	if cap(pl) < len(pool.pool) {
		ups := make([]*upstream, cap(pl)+10)
		for k, v := range pl {
			ups[k] = v
		}
		return ups
	}
	return pl
}

func (pool *streampool) pickupstreams() []*upstream {
	pool.waitforalive()

	// pick udp and tcp equally
	pool.mux.RLock()
	defer pool.mux.RUnlock()

	// pick one of each

	rn := int(rand.Int31())

	switch {
	case pool.tcplen > 0 && pool.udplen > 0:
		// pick one of each
		return []*upstream{
			pool.udpool[rn%pool.udplen],
			pool.tcpool[rn%pool.tcplen],
			pool.udpool[(rn+1)%pool.udplen],
		}
	case pool.tcplen == 0 || pool.udplen == 0:
		// pick 1-2 alived
		return []*upstream{
			pool.alived[rn%pool.alvlen],
			pool.alived[(rn+1)%pool.alvlen],
		}
	}
	logrus.Warnln("no upstream avalible for pick")
	return nil

}

func (pool *streampool) waitforalive() {
	pool.L.Lock()
	for !pool.alive {
		pool.Wait()
	}
	pool.L.Unlock()
}

func (pool *streampool) updateloop() {
	for {
		pool.L.Lock()
		for !pool.updatealive() {
			pool.Wait()
		}
		pool.L.Unlock()
		pool.Broadcast()
{'fraction_non_alphanumeric': 0.15967246673490276, 'fraction_numerical': 0.00831627430910952, 'mean_word_length': 2.932092555331992, 'pattern_counts': {'":': 20, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7609603', 'n_tokens_mistral': 3490, 'n_tokens_neox': 3198, 'n_words': 718}