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

package barnacle

import (
	"fmt"
	"sync"
	"time"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/vt/topo"
)

var idGen sync2.AtomicInt64

type VTConn struct {
	Id               int64
	balancerMap      *BalancerMap
	tabletType       topo.TabletType
	retryDelay       time.Duration
	retryCount       int
	shardConns       map[string]*ShardConn
	transactionId    int64
	transactionConns []*ShardConn
}

func NewVTConn(blm *BalancerMap, tabletType topo.TabletType, retryDelay time.Duration, retryCount int) *VTConn {
	return &VTConn{
		Id:          idGen.Add(1),
		balancerMap: blm,
		tabletType:  tabletType,
		retryDelay:  retryDelay,
		retryCount:  retryCount,
		shardConns:  make(map[string]*ShardConn),
	}
}

// Close closes the underlying ShardConn connections.
func (vtc *VTConn) Close() error {
	if vtc.shardConns == nil {
		return nil
	}
	for _, v := range vtc.shardConns {
		v.Close()
	}
	vtc.shardConns = nil
	vtc.balancerMap = nil
	return nil
}

func (vtc *VTConn) getConnection(keyspace, shard string) (*ShardConn, error) {
	key := fmt.Sprintf("%s.%s.%s", keyspace, vtc.tabletType, shard)
	sdc, ok := vtc.shardConns[key]
	if !ok {
		sdc = NewShardConn(vtc.balancerMap, keyspace, shard, vtc.tabletType, vtc.retryDelay, vtc.retryCount)
		vtc.shardConns[key] = sdc
	}
	if vtc.transactionId != 0 {
		if sdc.InTransaction() {
			return sdc, nil
		}
		if err := sdc.Begin(); err != nil {
			return nil, err
		}
		vtc.transactionConns = append(vtc.transactionConns, sdc)
	}
	return sdc, nil
}

func (vtc *VTConn) appendResult(qr, innerqr *mproto.QueryResult) {
	if qr.Fields == nil {
		qr.Fields = innerqr.Fields
	}
	qr.RowsAffected += innerqr.RowsAffected
	if innerqr.InsertId != 0 {
		qr.InsertId = innerqr.InsertId
	}
	qr.Rows = append(qr.Rows, innerqr.Rows...)
}

// ExecDirect executes a non-streaming query on the specified shards.
func (vtc *VTConn) ExecDirect(query string, bindVars map[string]interface{}, keyspace string, shards []string) (qr *mproto.QueryResult, err error) {
	qr = &mproto.QueryResult{}
	if len(shards) == 0 {
		return qr, nil
	}
	resultChan := make(chan *mproto.QueryResult, len(shards))
	errorChan := make(chan error, len(shards))
	var wg sync.WaitGroup
	wg.Add(len(shards))
	for _, shard := range shards {
		go func() {
			defer wg.Done()
			sdc, err := vtc.getConnection(keyspace, shard)
			if err != nil {
				errorChan <- err
				return
			}
			innerqr, err := sdc.ExecDirect(query, bindVars)
			if err != nil {
				errorChan <- err
				return<commit_msg>vtconn fastpath optimization<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package barnacle

import (
	"fmt"
	"sync"
	"time"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/vt/topo"
)

var idGen sync2.AtomicInt64

type VTConn struct {
	Id               int64
	balancerMap      *BalancerMap
	tabletType       topo.TabletType
	retryDelay       time.Duration
	retryCount       int
	shardConns       map[string]*ShardConn
	transactionId    int64
	transactionConns []*ShardConn
}

func NewVTConn(blm *BalancerMap, tabletType topo.TabletType, retryDelay time.Duration, retryCount int) *VTConn {
	return &VTConn{
		Id:          idGen.Add(1),
		balancerMap: blm,
		tabletType:  tabletType,
		retryDelay:  retryDelay,
		retryCount:  retryCount,
		shardConns:  make(map[string]*ShardConn),
	}
}

// Close closes the underlying ShardConn connections.
func (vtc *VTConn) Close() error {
	if vtc.shardConns == nil {
		return nil
	}
	for _, v := range vtc.shardConns {
		v.Close()
	}
	vtc.shardConns = nil
	vtc.balancerMap = nil
	return nil
}

func (vtc *VTConn) getConnection(keyspace, shard string) (*ShardConn, error) {
	key := fmt.Sprintf("%s.%s.%s", keyspace, vtc.tabletType, shard)
	sdc, ok := vtc.shardConns[key]
	if !ok {
		sdc = NewShardConn(vtc.balancerMap, keyspace, shard, vtc.tabletType, vtc.retryDelay, vtc.retryCount)
		vtc.shardConns[key] = sdc
	}
	if vtc.transactionId != 0 {
		if sdc.InTransaction() {
			return sdc, nil
		}
		if err := sdc.Begin(); err != nil {
			return nil, err
		}
		vtc.transactionConns = append(vtc.transactionConns, sdc)
	}
	return sdc, nil
}

func (vtc *VTConn) execOnShard(query string, bindVars map[string]interface{}, keyspace string, shard string) (qr *mproto.QueryResult, err error) {
	sdc, err := vtc.getConnection(keyspace, shard)
	if err != nil {
		return nil, err
	}
	qr, err = sdc.ExecDirect(query, bindVars)
	if err != nil {
		return nil, err
	}
	return qr, nil
}

func (vtc *VTConn) appendResult(qr, innerqr *mproto.QueryResult) {
	if qr.Fields == nil {
		qr.Fields = innerqr.Fields
	}
	qr.RowsAffected += innerqr.RowsAffected
	if innerqr.InsertId != 0 {
		qr.InsertId = innerqr.InsertId
	}
	qr.Rows = append(qr.Rows, innerqr.Rows...)
}

// ExecDirect executes a non-streaming query on the specified shards.
func (vtc *VTConn) ExecDirect(query string, bindVars map[string]interface{}, keyspace string, shards []string) (qr *mproto.QueryResult, err error) {
	switch len(shards) {
	case 0:
		return nil, nil
	case 1:
		return vtc.execOnShard(query, bindVars, keyspace, shards[0])
	}
	qr = new(mproto.QueryResult)
	resultChan := make(chan *mproto.QueryResult, len(shards))
	errorChan := make(chan error, len(shards))
	var wg sync.WaitGroup
	wg.Add(len(shards))
	for _, shard := range shards {
		go func() {
			defer wg.Done()
			innerqr, err := vtc.execOnShard(query, bindVars, keyspace, shard)
			if err != nil {
				errorChan <- err
				return
{'fraction_non_alphanumeric': 0.12154031287605295, 'fraction_numerical': 0.005844937252879491, 'mean_word_length': 3.6768488745980705, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 3, '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': '3709203', 'n_tokens_mistral': 2329, 'n_tokens_neox': 2179, 'n_words': 543}