<commit_before>package wrap

import (
	"io"
	"net/rpc"
	"sync"
	"time"

	"github.com/go-playground/log"
	"github.com/go-playground/wave/client"
)

// RetryReconnectEndpoint is interface needed to comply with wrapping
// with Rety and Reconnect logic.
type RetryReconnectEndpoint interface {
	client.Endpoint
	SetClient(*rpc.Client)
	NewClient() (*rpc.Client, error)
}

// RetryReconnect wraps the given RetryReconnectEndpoint endpoint and automatically
// handles logic to reconnect and retry
func RetryReconnect(endpoint RetryReconnectEndpoint, retryDuration time.Duration) (e client.Endpoint, err error) {

	rr := &retryReconnect{
		RetryReconnectEndpoint: endpoint,
		clientMutex:            new(sync.RWMutex),
		reconnectMutex:         new(sync.RWMutex),
		retryDuration:          retryDuration,
		isDisconnected:         true,
	}

	e = rr

	_, err = rr.NewClient()
	if err != nil {
		return
	}

	return
}

type retryReconnect struct {
	RetryReconnectEndpoint
	clientMutex    *sync.RWMutex
	reconnectMutex *sync.RWMutex
	retryDuration  time.Duration
	isDisconnected bool // defaults to true
}

var _ client.Endpoint = &retryReconnect{}
var _ RetryReconnectEndpoint = &retryReconnect{}

func (r *retryReconnect) NewClient() (c *rpc.Client, err error) {

	r.reconnectMutex.Lock()
	r.isDisconnected = true
	r.reconnectMutex.Unlock()

	// a bunch of calls could have gotten here if server was busy, lets double check we're still disconnected
	// right after lock is released
	r.clientMutex.Lock()
	defer r.clientMutex.Unlock()

	// check if still disconnected
	r.reconnectMutex.RLock()
	if !r.isDisconnected {
		r.reconnectMutex.RUnlock()
		return nil, err
	}
	r.reconnectMutex.RUnlock()

	for i := 0; i < 3; i++ {
		c, err = r.RetryReconnectEndpoint.NewClient()
		if err != nil {
			time.Sleep(r.retryDuration)
			continue
		}

		r.SetClient(c)

		r.reconnectMutex.Lock()
		r.isDisconnected = false
		r.reconnectMutex.Unlock()

		return
	}

	log.WithFields(log.F("err", err)).Alert("RPC Connection could not be established/reestablished")

	go func() {

		r.clientMutex.Lock()
		defer r.clientMutex.Unlock()

		for {

			var client *rpc.Client
			var err2 error

			// time.Sleep(r.retryDuration)
			client, err2 = r.RetryReconnectEndpoint.NewClient()
			if err2 != nil {
				time.Sleep(r.retryDuration)
				continue
			}

			r.SetClient(client)

			r.reconnectMutex.Lock()
			r.isDisconnected = false
			r.reconnectMutex.Unlock()

			break
		}
	}()

	err = rpc.ErrShutdown
	return
}
<commit_msg>re-add commented out sleep<commit_after>package wrap

import (
	"io"
	"net/rpc"
	"sync"
	"time"

	"github.com/go-playground/log"
	"github.com/go-playground/wave/client"
)

// RetryReconnectEndpoint is interface needed to comply with wrapping
// with Rety and Reconnect logic.
type RetryReconnectEndpoint interface {
	client.Endpoint
	SetClient(*rpc.Client)
	NewClient() (*rpc.Client, error)
}

// RetryReconnect wraps the given RetryReconnectEndpoint endpoint and automatically
// handles logic to reconnect and retry
func RetryReconnect(endpoint RetryReconnectEndpoint, retryDuration time.Duration) (e client.Endpoint, err error) {

	rr := &retryReconnect{
		RetryReconnectEndpoint: endpoint,
		clientMutex:            new(sync.RWMutex),
		reconnectMutex:         new(sync.RWMutex),
		retryDuration:          retryDuration,
		isDisconnected:         true,
	}

	e = rr

	_, err = rr.NewClient()
	if err != nil {
		return
	}

	return
}

type retryReconnect struct {
	RetryReconnectEndpoint
	clientMutex    *sync.RWMutex
	reconnectMutex *sync.RWMutex
	retryDuration  time.Duration
	isDisconnected bool // defaults to true
}

var _ client.Endpoint = &retryReconnect{}
var _ RetryReconnectEndpoint = &retryReconnect{}

func (r *retryReconnect) NewClient() (c *rpc.Client, err error) {

	r.reconnectMutex.Lock()
	r.isDisconnected = true
	r.reconnectMutex.Unlock()

	// a bunch of calls could have gotten here if server was busy, lets double check we're still disconnected
	// right after lock is released
	r.clientMutex.Lock()
	defer r.clientMutex.Unlock()

	// check if still disconnected
	r.reconnectMutex.RLock()
	if !r.isDisconnected {
		r.reconnectMutex.RUnlock()
		return nil, err
	}
	r.reconnectMutex.RUnlock()

	for i := 0; i < 3; i++ {
		c, err = r.RetryReconnectEndpoint.NewClient()
		if err != nil {
			time.Sleep(r.retryDuration)
			continue
		}

		r.SetClient(c)

		r.reconnectMutex.Lock()
		r.isDisconnected = false
		r.reconnectMutex.Unlock()

		return
	}

	log.WithFields(log.F("err", err)).Alert("RPC Connection could not be established/reestablished")

	go func() {

		r.clientMutex.Lock()
		defer r.clientMutex.Unlock()

		for {

			var client *rpc.Client
			var err2 error

			time.Sleep(r.retryDuration)
			client, err2 = r.RetryReconnectEndpoint.NewClient()
			if err2 != nil {
				time.Sleep(r.retryDuration)
				continue
			}

			r.SetClient(client)

			r.reconnectMutex.Lock()
			r.isDisconnected = false
			r.reconnectMutex.Unlock()

			break
		}
	}()

	err = rpc.ErrShutdown
	return
}

{'fraction_non_alphanumeric': 0.1069144684252598, 'fraction_numerical': 0.0019984012789768186, 'mean_word_length': 3.6256931608133085, 'pattern_counts': {'":': 0, '<': 5, '<?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': '909365', 'n_tokens_mistral': 1811, 'n_tokens_neox': 1716, 'n_words': 407}