<commit_before>package etcd

import (
	"context"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	etcdv3 "github.com/coreos/etcd/clientv3"
	etcdv3sync "github.com/coreos/etcd/clientv3/concurrency"
	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/state"
	"github.com/hashicorp/terraform/state/remote"
)

const (
	lockAcquireTimeout = 2 * time.Second
	lockInfoSuffix     = ".lockinfo"
	lockSuffix         = ".lock"
)

// RemoteClient is a remote client that will store data in etcd.
type RemoteClient struct {
	Client *etcdv3.Client
	DoLock bool
	Key    string

	etcdMutex   *etcdv3sync.Mutex
	etcdSession *etcdv3sync.Session
	info        *state.LockInfo
	mu          sync.Mutex
	modRevision int64
}

func (c *RemoteClient) Get() (*remote.Payload, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	res, err := c.Client.KV.Get(context.TODO(), c.Key)
	if err != nil {
		return nil, err
	}
	if res.Count == 0 {
		return nil, nil
	}
	if res.Count >= 2 {
		return nil, fmt.Errorf("Expected a single result but got %d.", res.Count)
	}

	c.modRevision = res.Kvs[0].ModRevision

	payload := res.Kvs[0].Value
	md5 := md5.Sum(payload)

	return &remote.Payload{
		Data: payload,
		MD5:  md5[:],
	}, nil
}

func (c *RemoteClient) Put(data []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	res, err := etcdv3.NewKV(c.Client).Txn(context.TODO()).If(
		etcdv3.Compare(etcdv3.ModRevision(c.Key), "=", c.modRevision),
	).Then(
		etcdv3.OpPut(c.Key, string(data)),
		etcdv3.OpGet(c.Key),
	).Commit()

	if err != nil {
		return err
	}
	if !res.Succeeded {
		return fmt.Errorf("The transaction did not succeed.")
	}
	if len(res.Responses) != 2 {
		return fmt.Errorf("Expected two responses but got %d.", len(res.Responses))
	}

	c.modRevision = res.Responses[1].GetResponseRange().Kvs[0].ModRevision
	return nil
}

func (c *RemoteClient) Delete() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	_, err := c.Client.KV.Delete(context.TODO(), c.Key)
	return err
}

func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.DoLock {
		return "", nil
	}

	c.info = info<commit_msg>Make sure we don't relock.<commit_after>package etcd

import (
	"context"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	etcdv3 "github.com/coreos/etcd/clientv3"
	etcdv3sync "github.com/coreos/etcd/clientv3/concurrency"
	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/state"
	"github.com/hashicorp/terraform/state/remote"
)

const (
	lockAcquireTimeout = 2 * time.Second
	lockInfoSuffix     = ".lockinfo"
	lockSuffix         = ".lock"
)

// RemoteClient is a remote client that will store data in etcd.
type RemoteClient struct {
	Client *etcdv3.Client
	DoLock bool
	Key    string

	etcdMutex   *etcdv3sync.Mutex
	etcdSession *etcdv3sync.Session
	info        *state.LockInfo
	mu          sync.Mutex
	modRevision int64
}

func (c *RemoteClient) Get() (*remote.Payload, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	res, err := c.Client.KV.Get(context.TODO(), c.Key)
	if err != nil {
		return nil, err
	}
	if res.Count == 0 {
		return nil, nil
	}
	if res.Count >= 2 {
		return nil, fmt.Errorf("Expected a single result but got %d.", res.Count)
	}

	c.modRevision = res.Kvs[0].ModRevision

	payload := res.Kvs[0].Value
	md5 := md5.Sum(payload)

	return &remote.Payload{
		Data: payload,
		MD5:  md5[:],
	}, nil
}

func (c *RemoteClient) Put(data []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	res, err := etcdv3.NewKV(c.Client).Txn(context.TODO()).If(
		etcdv3.Compare(etcdv3.ModRevision(c.Key), "=", c.modRevision),
	).Then(
		etcdv3.OpPut(c.Key, string(data)),
		etcdv3.OpGet(c.Key),
	).Commit()

	if err != nil {
		return err
	}
	if !res.Succeeded {
		return fmt.Errorf("The transaction did not succeed.")
	}
	if len(res.Responses) != 2 {
		return fmt.Errorf("Expected two responses but got %d.", len(res.Responses))
	}

	c.modRevision = res.Responses[1].GetResponseRange().Kvs[0].ModRevision
	return nil
}

func (c *RemoteClient) Delete() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	_, err := c.Client.KV.Delete(context.TODO(), c.Key)
	return err
}

func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.DoLock {
		return "", nil
	}
	if c.etcdSession != nil {
		return "", fmt.Errorf("state %q already locked", c.Key)
	}

	c.info = info
{'fraction_non_alphanumeric': 0.15652972228597659, 'fraction_numerical': 0.012393848978655038, 'mean_word_length': 3.5301455301455302, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '7572762', 'n_tokens_mistral': 1815, 'n_tokens_neox': 1750, 'n_words': 374}