<commit_before>package client

import (
	"context"
	"time"

	"koding/klient/machine/index"
)

// DynamicClientFunc is an adapter that allows to dynamically provide clients.
type DynamicClientFunc func() (Client, error)

// Supervised is a decorator type for Clients which will wait for valid client
// if the current client is Disconnected. This type is meant to handle temporary
// network issues or cases when underlying client is not set up yet. Eg. when
// program started but haven't made connection to remote machine yet.
type Supervised struct {
	dcf     DynamicClientFunc
	timeout time.Duration
}

// NewSupervised creates a new Supervised client instance.
func NewSupervised(dcf DynamicClientFunc, timeout time.Duration) *Supervised {
	return &Supervised{
		timeout: timeout,
		dcf:     dcf,
	}
}

// CurrentUser calls registered Client's CurrentUser method and returns its
// result if it's not produced by Disconnected client. If it is, this function
// will wait until valid client is available or timeout is reached.
func (s *Supervised) CurrentUser() (user string, err error) {
	if e := s.call(func(c Client) error {
		user, err = c.CurrentUser()
		return err
	}); e != nil {
		return "", e
	}

	return
}

// SSHAddKeys calls registered Client's SSHAddKeys method and returns its result
// if it's not produced by Disconnected client. If it is, this function will
// wait until valid client is available or timeout is reached.
func (s *Supervised) SSHAddKeys(username string, keys ...string) (err error) {
	if e := s.call(func(c Client) error {
		err = c.SSHAddKeys(username, keys...)
		return err
	}); e != nil {
		return e
	}

	return
}

// MountHeadIndex calls registered Client's MountHeadIndex method and returns
// its result if it's not produced by Disconnected client. If it is, this
// function will wait until valid client is available or timeout is reached.
func (s *Supervised) MountHeadIndex(path string) (absPath string, count int, diskSize int64, err error) {
	if e := s.call(func(c Client) error {
		absPath, count, diskSize, err = c.MountHeadIndex(path)
		return err
	}); e != nil {
		return "", 0, 0, e
	}

	return
}

// MountGetIndex calls registered Client's MountGetIndex method and returns its
// result if it's not produced by Disconnected client. If it is, this function
// will wait until valid client is available or timeout is reached.
func (s *Supervised) MountGetIndex(path string) (idx *index.Index, err error) {
	if e := s.call(func(c Client) error {
		idx, err = c.MountGetIndex(path)
		return err
	}); e != nil {
		return nil, e
	}

	return
}

// Context calls registered Client's Context method and returns its result. If
// there is an error during client retrieving, this function will return
// canceled context.
func (s *Supervised) Context() context.Context {
	c, err := s.dcf()
	if err == nil {
		return c.Context()
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	return ctx
}

func (s *Supervised) call(f func(Client) error) error {
	c, err := s.dcf()
	if err != nil {
		return err
	}

	ctx := c.Context()
	if err = f(c); err == ErrDisconnected {
		timer := time.NewTimer(s.timeout)
		select {
		case <-ctx.Done():
			// Client changed. Use a new one.
			timer.Stop()
			c, err = s.dcf()
			if err != nil {
				return err
			}
			return f(c)
		case <-timer.C:
			// Client is still disconnected. Return it as is.
			return nil
		}
	}

	return err
}<commit_msg>klient/machine: use context with timeout to handle client timeouts<commit_after>package client

import (
	"context"
	"time"

	"koding/klient/machine/index"
)

// DynamicClientFunc is an adapter that allows to dynamically provide clients.
type DynamicClientFunc func() (Client, error)

// Supervised is a decorator type for Clients which will wait for valid client
// if the current client is Disconnected. This type is meant to handle temporary
// network issues or cases when underlying client is not set up yet. Eg. when
// program started but haven't made connection to remote machine yet.
type Supervised struct {
	dcf     DynamicClientFunc
	timeout time.Duration
}

// NewSupervised creates a new Supervised client instance.
func NewSupervised(dcf DynamicClientFunc, timeout time.Duration) *Supervised {
	return &Supervised{
		timeout: timeout,
		dcf:     dcf,
	}
}

// CurrentUser calls registered Client's CurrentUser method and returns its
// result if it's not produced by Disconnected client. If it is, this function
// will wait until valid client is available or timeout is reached.
func (s *Supervised) CurrentUser() (user string, err error) {
	if e := s.call(func(c Client) error {
		user, err = c.CurrentUser()
		return err
	}); e != nil {
		return "", e
	}

	return
}

// SSHAddKeys calls registered Client's SSHAddKeys method and returns its result
// if it's not produced by Disconnected client. If it is, this function will
// wait until valid client is available or timeout is reached.
func (s *Supervised) SSHAddKeys(username string, keys ...string) (err error) {
	if e := s.call(func(c Client) error {
		err = c.SSHAddKeys(username, keys...)
		return err
	}); e != nil {
		return e
	}

	return
}

// MountHeadIndex calls registered Client's MountHeadIndex method and returns
// its result if it's not produced by Disconnected client. If it is, this
// function will wait until valid client is available or timeout is reached.
func (s *Supervised) MountHeadIndex(path string) (absPath string, count int, diskSize int64, err error) {
	if e := s.call(func(c Client) error {
		absPath, count, diskSize, err = c.MountHeadIndex(path)
		return err
	}); e != nil {
		return "", 0, 0, e
	}

	return
}

// MountGetIndex calls registered Client's MountGetIndex method and returns its
// result if it's not produced by Disconnected client. If it is, this function
// will wait until valid client is available or timeout is reached.
func (s *Supervised) MountGetIndex(path string) (idx *index.Index, err error) {
	if e := s.call(func(c Client) error {
		idx, err = c.MountGetIndex(path)
		return err
	}); e != nil {
		return nil, e
	}

	return
}

// Context calls registered Client's Context method and returns its result. If
// there is an error during client retrieving, this function will return
// canceled context.
func (s *Supervised) Context() context.Context {
	c, err := s.dcf()
	if err != nil {
		ctx, cancel := context.WithCancel(context.Background())
		cancel()
		return ctx
	}

	return c.Context()
}

func (s *Supervised) call(f func(Client) error) error {
	c, err := s.dcf()
	if err != nil {
		return err
	}

	ctx := c.Context()
	if err = f(c); err != ErrDisconnected {
		return err
	}

	// Wait for new client.
	ctx, cancel := context.WithTimeout(ctx, s.timeout)
	defer cancel()

	if _ = <-ctx.Done(); ctx.Err() == context.DeadlineExceeded {
		// Client is still disconnected. Return it as is.
		return nil
	}

	// Previous context was canceled. This means that the client changed.
	c, err = s.dcf()
	if err != nil {
		return err
	}

	return f(c)
}
{'fraction_non_alphanumeric': 0.10102593331433457, 'fraction_numerical': 0.001139925904816187, 'mean_word_length': 3.5875816993464054, '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': '7052580', 'n_tokens_mistral': 2214, 'n_tokens_neox': 2076, 'n_words': 887}