<commit_before>package mux

import (
	"context"
	"sync"
	"time"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/ray"
)

const (
	maxParallel = 8
	maxTotal    = 128
)

type manager interface {
	remove(id uint16)
}

type session struct {
	sync.Mutex
	input          ray.InputStream
	output         ray.OutputStream
	parent         manager
	id             uint16
	uplinkClosed   bool
	downlinkClosed bool
}

func (s *session) checkAndRemove() {
	s.Lock()
	if s.uplinkClosed && s.downlinkClosed {
		s.parent.remove(s.id)
	}
	s.Unlock()
}

func (s *session) closeUplink() {
	s.Lock()
	s.uplinkClosed = true
	s.Unlock()
	s.checkAndRemove()
}

func (s *session) closeDownlink() {
	s.Lock()
	s.downlinkClosed = true
	s.Unlock()
	s.checkAndRemove()
}

type ClientManager struct {
	access  sync.Mutex
	clients []*Client
	proxy   proxy.Outbound
	dialer  proxy.Dialer
}

func NewClientManager(p proxy.Outbound, d proxy.Dialer) *ClientManager {
	return &ClientManager{
		proxy:  p,
		dialer: d,
	}
}

func (m *ClientManager) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) error {
	m.access.Lock()
	defer m.access.Unlock()

	for _, client := range m.clients {
		if client.Dispatch(ctx, outboundRay) {
			return nil
		}
	}

	client, err := NewClient(m.proxy, m.dialer, m)
	if err != nil {
		return errors.Base(err).Message("Proxyman|Mux|ClientManager: Failed to create client.")
	}
	m.clients = append(m.clients, client)
	client.Dispatch(ctx, outboundRay)
	return nil
}

func (m *ClientManager) onClientFinish() {
	m.access.Lock()
	defer m.access.Unlock()

	if len(m.clients) < 10 {
		return
	}

	activeClients := make([]*Client, 0, len(m.clients))

	for _, client := range m.clients {
		if !client.Closed() {
			activeClients = append(activeClients, client)
		}
	}
	m.clients = activeClients
}

type Client struct {
	access     sync.RWMutex
	count      uint16
	sessions   map[uint16]*session
	inboundRay ray.InboundRay
	ctx        context.Context
	cancel     context.CancelFunc
	manager    *ClientManager
}

var muxCoolDestination = net.TCPDestination(net.DomainAddress("v1.mux.cool"), net.Port(9527))

func NewClient(p proxy.Outbound, dialer proxy.Dialer, m *ClientManager) (*Client, error) {
	ctx, cancel := context.WithCancel(context.Background())
	ctx = proxy.ContextWithTarget(ctx, muxCoolDestination)
	pipe := ray.NewRay(ctx)
	go p.Process(ctx, pipe, dialer)
	c := &Client{
		sessions:   make(map[uint16]*session, 256),
		inboundRay: pipe,
		ctx:        ctx,
		cancel:     cancel,
		manager:    m,
		count:      0,
	}
	go c.fetchOutput()
	return c, nil
}

func (m *Client) remove(id uint16) {
	m.access.Lock()
	defer m.access.Unlock()

	delete(m.sessions, id)

	if len(m.sessions) == 0 {
		m.cancel()
		m.inboundRay.InboundInput().Close()
		go m.manager.onClientFinish()
	}
}

func (m *Client) Closed() bool {
	select {
	case <-m.ctx.Done():
		return true
	default:
		return false
	}
}

func fetchInput(ctx context.Context, s *session, output buf.Writer) {
	dest, _ := proxy.TargetFromContext(ctx)
	writer := &Writer{
		dest:   dest,
		id:     s.id,
		writer: output,
	}
	defer writer.Close()
	defer s.closeUplink()

	log.Info("Proxyman|Mux|Client: Dispatching request to ", dest)
	data, _ := s.input.ReadTimeout(time.Millisecond * 500)
	if data != nil {
		if err := writer.Write(data); err != nil {
			log.Info("Proxyman|Mux|Client: Failed to write first payload: ", err)
			return
		}
	}
	_, timer := signal.CancelAfterInactivity(ctx, time.Minute*5)
	if err := buf.PipeUntilEOF(timer, s.input, writer); err != nil {
		log.Info("Proxyman|Mux|Client: Failed to fetch all input: ", err)
	}
}

func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool {
	m.access.Lock()
	defer m.access.Unlock()

	if len(m.sessions) >= maxParallel {
		return false
	}

	if m.count >= maxTotal {
		return false
	}

	select {<commit_msg>simplify mux code<commit_after>package mux

import (
	"context"
	"sync"
	"time"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/ray"
)

const (
	maxParallel = 8
	maxTotal    = 128
)

type manager interface {
	remove(id uint16)
}

type session struct {
	sync.Mutex
	input          ray.InputStream
	output         ray.OutputStream
	parent         manager
	id             uint16
	uplinkClosed   bool
	downlinkClosed bool
}

func (s *session) closeUplink() {
	var allDone bool
	s.Lock()
	s.uplinkClosed = true
	allDone = s.uplinkClosed && s.downlinkClosed
	s.Unlock()
	if allDone {
		s.parent.remove(s.id)
	}
}

func (s *session) closeDownlink() {
	var allDone bool
	s.Lock()
	s.downlinkClosed = true
	allDone = s.uplinkClosed && s.downlinkClosed
	s.Unlock()
	if allDone {
		s.parent.remove(s.id)
	}
}

type ClientManager struct {
	access  sync.Mutex
	clients []*Client
	proxy   proxy.Outbound
	dialer  proxy.Dialer
}

func NewClientManager(p proxy.Outbound, d proxy.Dialer) *ClientManager {
	return &ClientManager{
		proxy:  p,
		dialer: d,
	}
}

func (m *ClientManager) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) error {
	m.access.Lock()
	defer m.access.Unlock()

	for _, client := range m.clients {
		if client.Dispatch(ctx, outboundRay) {
			return nil
		}
	}

	client, err := NewClient(m.proxy, m.dialer, m)
	if err != nil {
		return errors.Base(err).Message("Proxyman|Mux|ClientManager: Failed to create client.")
	}
	m.clients = append(m.clients, client)
	client.Dispatch(ctx, outboundRay)
	return nil
}

func (m *ClientManager) onClientFinish() {
	m.access.Lock()
	defer m.access.Unlock()

	if len(m.clients) < 10 {
		return
	}

	activeClients := make([]*Client, 0, len(m.clients))

	for _, client := range m.clients {
		if !client.Closed() {
			activeClients = append(activeClients, client)
		}
	}
	m.clients = activeClients
}

type Client struct {
	access     sync.RWMutex
	count      uint16
	sessions   map[uint16]*session
	inboundRay ray.InboundRay
	ctx        context.Context
	cancel     context.CancelFunc
	manager    *ClientManager
}

var muxCoolDestination = net.TCPDestination(net.DomainAddress("v1.mux.cool"), net.Port(9527))

func NewClient(p proxy.Outbound, dialer proxy.Dialer, m *ClientManager) (*Client, error) {
	ctx, cancel := context.WithCancel(context.Background())
	ctx = proxy.ContextWithTarget(ctx, muxCoolDestination)
	pipe := ray.NewRay(ctx)
	go p.Process(ctx, pipe, dialer)
	c := &Client{
		sessions:   make(map[uint16]*session, 256),
		inboundRay: pipe,
		ctx:        ctx,
		cancel:     cancel,
		manager:    m,
		count:      0,
	}
	go c.fetchOutput()
	return c, nil
}

func (m *Client) remove(id uint16) {
	m.access.Lock()
	defer m.access.Unlock()

	delete(m.sessions, id)

	if len(m.sessions) == 0 {
		m.cancel()
		m.inboundRay.InboundInput().Close()
		go m.manager.onClientFinish()
	}
}

func (m *Client) Closed() bool {
	select {
	case <-m.ctx.Done():
		return true
	default:
		return false
	}
}

func fetchInput(ctx context.Context, s *session, output buf.Writer) {
	dest, _ := proxy.TargetFromContext(ctx)
	writer := &Writer{
		dest:   dest,
		id:     s.id,
		writer: output,
	}
	defer writer.Close()
	defer s.closeUplink()

	log.Info("Proxyman|Mux|Client: Dispatching request to ", dest)
	data, _ := s.input.ReadTimeout(time.Millisecond * 500)
	if data != nil {
		if err := writer.Write(data); err != nil {
			log.Info("Proxyman|Mux|Client: Failed to write first payload: ", err)
			return
		}
	}
	_, timer := signal.CancelAfterInactivity(ctx, time.Minute*30)
	if err := buf.PipeUntilEOF(timer, s.input, writer); err != nil {
		log.Info("Proxyman|Mux|Client: Failed to fetch all input: ", err)
	}
}

func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool {
	m.access.Lock()
	defer m.access.Unlock()

	if len(m.sessions) >= maxParallel {
		return false
	}

	if m.count >= maxTotal {
		return false
	}

	select {
{'fraction_non_alphanumeric': 0.1362464710936541, 'fraction_numerical': 0.01043328832699153, 'mean_word_length': 3.6085972850678734, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, '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': '7450777', 'n_tokens_mistral': 3254, 'n_tokens_neox': 3088, 'n_words': 678}