github.com/orofarne/hammy@v0.0.0-20130409105742-374fadfd6ecb/src/hammy/couchbase_state.go (about)

     1  package hammy
     2  
     3  import (
     4  	"fmt"
     5  	"encoding/json"
     6  	"encoding/binary"
     7  	"github.com/couchbaselabs/go-couchbase"
     8  	"github.com/dustin/gomemcached"
     9  	"github.com/dustin/gomemcached/client"
    10  )
    11  
    12  type CouchbaseStateKeeper struct {
    13  	Client *couchbase.Client
    14  	Pool *couchbase.Pool
    15  	Bucket *couchbase.Bucket
    16  	Ttl uint32
    17  }
    18  
    19  func NewCouchbaseStateKeeper(cfg Config) (*CouchbaseStateKeeper, error) {
    20  	sk := new(CouchbaseStateKeeper)
    21  
    22  	c, err := couchbase.Connect(cfg.CouchbaseStates.ConnectTo)
    23  	if err != nil {
    24  		return nil, err
    25  	}
    26  	sk.Client = &c
    27  
    28  	p, err := sk.Client.GetPool(cfg.CouchbaseStates.Pool)
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  	sk.Pool = &p
    33  
    34  	b, err := sk.Pool.GetBucket(cfg.CouchbaseStates.Bucket)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  	sk.Bucket = b
    39  
    40  	sk.Ttl = uint32(cfg.CouchbaseStates.Ttl)
    41  
    42  	return sk, nil
    43  }
    44  
    45  func (sk *CouchbaseStateKeeper) Get(key string) StateKeeperAnswer {
    46  	s := NewState()
    47  	var cas uint64
    48  	err := sk.Bucket.Gets(key, s, &cas)
    49  
    50  	if err == nil {
    51  		return StateKeeperAnswer{
    52  			State: *s,
    53  			Cas: &cas,
    54  			Err: nil,
    55  		}
    56  	} else {
    57  		return StateKeeperAnswer{
    58  			State: nil,
    59  			Cas: nil,
    60  			Err: err,
    61  		}
    62  	}
    63  	panic("?!!")
    64  }
    65  
    66  func (sk *CouchbaseStateKeeper) MGet(keys []string) (states map[string]StateKeeperAnswer) {
    67  	ans := sk.Bucket.GetBulk(keys)
    68  
    69  	states = make(map[string]StateKeeperAnswer)
    70  	for k, r := range ans {
    71  		switch r.Status {
    72  			case gomemcached.SUCCESS:
    73  				s := NewState()
    74  				err := json.Unmarshal(r.Body, s)
    75  				if err == nil {
    76  					states[k] = StateKeeperAnswer{
    77  						State: *s,
    78  						Cas: &r.Cas,
    79  						Err: nil,
    80  					}
    81  				} else {
    82  					states[k] = StateKeeperAnswer{
    83  						State: nil,
    84  						Cas: nil,
    85  						Err: err,
    86  					}
    87  				}
    88  			case gomemcached.KEY_ENOENT:
    89  				states[k] = StateKeeperAnswer{
    90  					State: *NewState(),
    91  					Cas: nil,
    92  					Err: nil,
    93  				}
    94  			default:
    95  				states[k] = StateKeeperAnswer{
    96  					State: nil,
    97  					Cas: nil,
    98  					Err: fmt.Errorf("%s", r.Error()),
    99  				}
   100  		}
   101  	}
   102  
   103  	for _, k := range keys {
   104  		if _, found := states[k]; !found {
   105  			states[k] = StateKeeperAnswer{
   106  				State: *NewState(),
   107  				Cas: nil,
   108  				Err: nil,
   109  			}
   110  		}
   111  	}
   112  
   113  	return
   114  }
   115  
   116  func (sk *CouchbaseStateKeeper) Set(key string, data State, cas *uint64) (retry bool, err error) {
   117  	err = sk.Bucket.Do(key, func(mc *memcached.Client, vb uint16) (e error) {
   118  		buf, e := json.Marshal(data)
   119  		if e != nil {
   120  			return
   121  		}
   122  		req := &gomemcached.MCRequest{
   123  			Opcode: gomemcached.SET,
   124  			VBucket: vb,
   125  			Key: []byte(key),
   126  			Cas: 0,
   127  			Opaque: 0,
   128  			Extras: []byte{0, 0, 0, 0, 0, 0, 0, 0},
   129  			Body: buf,
   130  		}
   131  		if cas != nil {
   132  			req.Cas = *cas
   133  		}
   134  		binary.BigEndian.PutUint32(req.Extras[4:8], sk.Ttl)
   135  
   136  		resp, e := mc.Send(req)
   137  		if e != nil {
   138  			if resp != nil && resp.Status == gomemcached.KEY_EEXISTS {
   139  				e = nil
   140  				retry = true
   141  			}
   142  			return
   143  		}
   144  
   145  		switch resp.Status {
   146  			case gomemcached.KEY_EEXISTS:
   147  				retry = true
   148  				return
   149  			case gomemcached.SUCCESS:
   150  				return
   151  			default:
   152  				return fmt.Errorf("CAS operation failed: %v", resp.Error())
   153  		}
   154  		panic("?!!")
   155  	})
   156  	if err != nil {
   157  		return
   158  	}
   159  
   160  	return
   161  }