<commit_before>package gocbcore

import (
	"crypto/tls"
	"encoding/binary"
	"net"
	"time"
)

func (c *Agent) waitAndRetryOperation(req *memdQRequest) {
	if c.nmvRetryDelay == 0 {
		c.redispatchDirect(req)
	} else {
		time.AfterFunc(c.nmvRetryDelay, func() {
			c.redispatchDirect(req)
		})
	}
}

func (c *Agent) handleServerNmv(s *memdPipeline, req *memdQRequest, resp *memdResponse) {
	// Try to parse the value as a bucket configuration
	bk, err := parseConfig(resp.Value, s.Hostname())
	if err == nil {
		c.updateConfig(bk)
	}

	// Redirect it!  This may actually come back to this server, but I won't tell
	//   if you don't ;)
	c.waitAndRetryOperation(req)
}

func (c *Agent) handleServerDeath(s *memdPipeline) {
	// Check if we are shutting down, if so, simply notify the shutdown
	//   method that we are going away.
	routeData := c.routingInfo.get()
	if routeData == nil {
		c.shutdownWaitCh <- s
		return
	}

	// Refresh the routing data with the existing configuration, this has
	//   the effect of attempting to rebuild the dead server.
	c.updateConfig(nil)

	// TODO(brett19): We probably should actually try other ways of resolving
	//  the issue, like requesting a new configuration.
}

func appendFeatureCode(bytes []byte, feature HelloFeature) []byte {
	bytes = append(bytes, 0, 0)
	binary.BigEndian.PutUint16(bytes[len(bytes)-2:], uint16(FeatureSeqNo))
	return bytes
}

func (c *Agent) tryHello(pipeline *memdPipeline, deadline time.Time) error {
	var featuresBytes []byte

	if c.useMutationTokens {
		featuresBytes = appendFeatureCode(featuresBytes, FeatureSeqNo)
	}

	_, err := pipeline.ExecuteRequest(&memdQRequest{
		memdRequest: memdRequest{
			Magic:  ReqMagic,
			Opcode: CmdHello,
			Key:    []byte("gocb/" + GoCbVersionStr),
			Value:  featuresBytes,
		},
	}, deadline)

	return err
}

// Attempt to connect a server, this function must be called
//  in its own goroutine and will ensure that offline servers
//  are not spammed with connection attempts.
func (agent *Agent) connectServer(server *memdPipeline) {
	agent.serverFailuresLock.Lock()
	failureTime := agent.serverFailures[server.address]
	agent.serverFailuresLock.Unlock()

	if !failureTime.IsZero() {
		waitedTime := time.Since(failureTime)
		if waitedTime < agent.serverWaitTimeout {
			time.Sleep(agent.serverWaitTimeout - waitedTime)

			if !agent.checkPendingServer(server) {
				// Server is no longer pending.  Stop trying.
				return
			}
		}
	}

	err := agent.connectPipeline(server, time.Now().Add(agent.serverConnectTimeout))
	if err != nil {
		agent.serverFailuresLock.Lock()
		agent.serverFailures[server.address] = time.Now()
		agent.serverFailuresLock.Unlock()

		// Force a config update which will clear away this dead pending
		// server and create a new one to connect with.
		agent.updateConfig(nil)
		return
	}

	if !agent.activatePendingServer(server) {
		// If this is no longer a valid pending server, we should shut
		//   it down!
		server.Close()
	}
}

func (c *Agent) connectPipeline(pipeline *memdPipeline, deadline time.Time) error {
	logDebugf("Attempting to connect pipeline to %s", pipeline.address)

	// Copy the tls configuration since we need to provide the hostname for each
	// server that we connect to so that the certificate can be validated properly.
	var tlsConfig *tls.Config = nil
	if c.tlsConfig != nil {
		host, _, _ := net.SplitHostPort(pipeline.address)
		tlsConfig = &tls.Config{
			InsecureSkipVerify: c.tlsConfig.InsecureSkipVerify,
			RootCAs:            c.tlsConfig.RootCAs,
			ServerName:         host,
		}
	}

	memdConn, err := DialMemdConn(pipeline.address, tlsConfig, deadline)
	if err != nil {
		logDebugf("Failed to connect. %v", err)
		pipeline.lock.Lock()
		pipeline.isClosed = true
		pipeline.lock.Unlock()
		return err
	}

	logDebugf("Connected.")
	pipeline.conn = memdConn
	go pipeline.Run()

	logDebugf("Authenticating...")
	err = c.initFn(pipeline, deadline)
	if err != nil {
		logDebugf("Failed to authenticate. %v", err)
		memdConn.Close()
		return err
	}

	c.tryHello(pipeline, deadline)

	return nil
}

// Drains all the requests out of the queue for this server.  This must be
//   invoked only once this server no longer exists in the routing data or an
//   infinite loop will likely occur.
func (c *Agent) shutdownPipeline(s *memdPipeline) {
	s.Drain(func(req *memdQRequest) {
		c.redispatchDirect(req)
	})
}

func (agent *Agent) checkPendingServer(server *memdPipeline) bool {
	oldRouting := agent.routingInfo.get()
	if oldRouting == nil {
		return false
	}

	// Find the index of the pending server we want to swap
	var serverIdx int = -1
	for i, s := range oldRouting.pendingServers {
		if s == server {
			serverIdx = i
			break
		}
	}

	return serverIdx != -1
}
<commit_msg>GOCBC-124: Properly close connections after shutdown.<commit_after>package gocbcore

import (
	"crypto/tls"
	"encoding/binary"
	"net"
	"time"
)

func (c *Agent) waitAndRetryOperation(req *memdQRequest) {
	if c.nmvRetryDelay == 0 {
		c.redispatchDirect(req)
	} else {
		time.AfterFunc(c.nmvRetryDelay, func() {
			c.redispatchDirect(req)
		})
	}
}

func (c *Agent) handleServerNmv(s *memdPipeline, req *memdQRequest, resp *memdResponse) {
	// Try to parse the value as a bucket configuration
	bk, err := parseConfig(resp.Value, s.Hostname())
	if err == nil {
		c.updateConfig(bk)
	}

	// Redirect it!  This may actually come back to this server, but I won't tell
	//   if you don't ;)
	c.waitAndRetryOperation(req)
}

func (c *Agent) handleServerDeath(s *memdPipeline) {
	// Check if we are shutting down, if so, simply notify the shutdown
	//   method that we are going away.
	routeData := c.routingInfo.get()
	if routeData == nil {
		c.shutdownWaitCh <- s
		return
	}

	// Refresh the routing data with the existing configuration, this has
	//   the effect of attempting to rebuild the dead server.
	c.updateConfig(nil)

	// TODO(brett19): We probably should actually try other ways of resolving
	//  the issue, like requesting a new configuration.
}

func appendFeatureCode(bytes []byte, feature HelloFeature) []byte {
	bytes = append(bytes, 0, 0)
	binary.BigEndian.PutUint16(bytes[len(bytes)-2:], uint16(FeatureSeqNo))
	return bytes
}

func (c *Agent) tryHello(pipeline *memdPipeline, deadline time.Time) error {
	var featuresBytes []byte

	if c.useMutationTokens {
		featuresBytes = appendFeatureCode(featuresBytes, FeatureSeqNo)
	}

	_, err := pipeline.ExecuteRequest(&memdQRequest{
		memdRequest: memdRequest{
			Magic:  ReqMagic,
			Opcode: CmdHello,
			Key:    []byte("gocb/" + GoCbVersionStr),
			Value:  featuresBytes,
		},
	}, deadline)

	return err
}

// Attempt to connect a server, this function must be called
//  in its own goroutine and will ensure that offline servers
//  are not spammed with connection attempts.
func (agent *Agent) connectServer(server *memdPipeline) {
	agent.serverFailuresLock.Lock()
	failureTime := agent.serverFailures[server.address]
	agent.serverFailuresLock.Unlock()

	if !failureTime.IsZero() {
		waitedTime := time.Since(failureTime)
		if waitedTime < agent.serverWaitTimeout {
			time.Sleep(agent.serverWaitTimeout - waitedTime)

			if !agent.checkPendingServer(server) {
				// Server is no longer pending.  Stop trying.
				return
			}
		}
	}

	err := agent.connectPipeline(server, time.Now().Add(agent.serverConnectTimeout))
	if err != nil {
		agent.serverFailuresLock.Lock()
		agent.serverFailures[server.address] = time.Now()
		agent.serverFailuresLock.Unlock()

		// Force a config update which will clear away this dead pending
		// server and create a new one to connect with.
		agent.updateConfig(nil)
		return
	}

	if !agent.activatePendingServer(server) {
		// If this is no longer a valid pending server, we should shut
		//   it down!
		server.Close()
	}
}

func (c *Agent) connectPipeline(pipeline *memdPipeline, deadline time.Time) error {
	logDebugf("Attempting to connect pipeline to %s", pipeline.address)

	// Copy the tls configuration since we need to provide the hostname for each
	// server that we connect to so that the certificate can be validated properly.
	var tlsConfig *tls.Config = nil
	if c.tlsConfig != nil {
		host, _, _ := net.SplitHostPort(pipeline.address)
		tlsConfig = &tls.Config{
			InsecureSkipVerify: c.tlsConfig.InsecureSkipVerify,
			RootCAs:            c.tlsConfig.RootCAs,
			ServerName:         host,
		}
	}

	memdConn, err := DialMemdConn(pipeline.address, tlsConfig, deadline)
	if err != nil {
		logDebugf("Failed to connect. %v", err)
		pipeline.lock.Lock()
		pipeline.isClosed = true
		pipeline.lock.Unlock()
		return err
	}

	logDebugf("Connected.")
	pipeline.conn = memdConn
	go pipeline.Run()

	logDebugf("Authenticating...")
	err = c.initFn(pipeline, deadline)
	if err != nil {
		logDebugf("Failed to authenticate. %v", err)
		memdConn.Close()
		return err
	}

	c.tryHello(pipeline, deadline)

	return nil
}

// Drains all the requests out of the queue for this server.  This must be
//   invoked only once this server no longer exists in the routing data or an
//   infinite loop will likely occur.
func (c *Agent) shutdownPipeline(s *memdPipeline) {
	s.Drain(func(req *memdQRequest) {
		c.redispatchDirect(req)
	})

	s.Close()
}

func (agent *Agent) checkPendingServer(server *memdPipeline) bool {
	oldRouting := agent.routingInfo.get()
	if oldRouting == nil {
		return false
	}

	// Find the index of the pending server we want to swap
	var serverIdx int = -1
	for i, s := range oldRouting.pendingServers {
		if s == server {
			serverIdx = i
			break
		}
	}

	return serverIdx != -1
}

{'fraction_non_alphanumeric': 0.10803440318858822, 'fraction_numerical': 0.0028319697923222154, 'mean_word_length': 3.791457286432161, 'pattern_counts': {'":': 0, '<': 7, '<?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': '7186815', 'n_tokens_mistral': 3309, 'n_tokens_neox': 3122, 'n_words': 936}