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

     1  package hammy
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"encoding/binary"
     7  	"github.com/ugorji/go-msgpack"
     8  	"github.com/couchbaselabs/go-couchbase"
     9  	"github.com/dustin/gomemcached"
    10  	"github.com/dustin/gomemcached/client"
    11  )
    12  
    13  // Saves historical data to write chache (based on couchbase)
    14  type CouchbaseSaver struct {
    15  	client *couchbase.Client
    16  	pool *couchbase.Pool
    17  	bucket *couchbase.Bucket
    18  	dataChan chan *IncomingData
    19  	Ttl uint32
    20  }
    21  
    22  // Create new saver
    23  func NewCouchbaseSaver(cfg Config) (*CouchbaseSaver, error) {
    24  	s := new(CouchbaseSaver)
    25  
    26  	c, err := couchbase.Connect(cfg.CouchbaseSaver.ConnectTo)
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  	s.client = &c
    31  
    32  	p, err := s.client.GetPool(cfg.CouchbaseSaver.Pool)
    33  	if err != nil {
    34  		return nil, err
    35  	}
    36  	s.pool = &p
    37  
    38  	b, err := s.pool.GetBucket(cfg.CouchbaseSaver.Bucket)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	s.bucket = b
    43  
    44  	s.Ttl = uint32(cfg.CouchbaseSaver.Ttl)
    45  
    46  	// Process queue
    47  	s.dataChan = make(chan *IncomingData, cfg.CouchbaseSaver.QueueSize)
    48  	// Workers...
    49  	for i := uint(0); i < cfg.CouchbaseSaver.SavePoolSize; i++ {
    50  		go s.worker()
    51  	}
    52  
    53  	return s, nil
    54  }
    55  
    56  // Enqueue data for saving
    57  func (s *CouchbaseSaver) Push(data *IncomingData) {
    58  	s.dataChan <- data
    59  }
    60  
    61  const CouchbaseDataBucketQuantum = 7200 // 2 hours
    62  
    63  func CouchbaseSaverBucketKey(hostKey string, itemKey string, timestamp uint64) string {
    64  	var bucketId uint64
    65  	bucketId = timestamp / CouchbaseDataBucketQuantum
    66  	return fmt.Sprintf("%s$%s$%d", hostKey, itemKey, bucketId)
    67  }
    68  
    69  func (s *CouchbaseSaver) worker() {
    70  	for data := range s.dataChan {
    71  		for hostK, hostV := range *data {
    72  			for itemK, itemV := range hostV {
    73  				for _, v := range itemV {
    74  					err := s.saveItem(hostK, itemK, v)
    75  					if err != nil {
    76  						log.Printf("saveItem error: %v", err)
    77  					}
    78  				}
    79  			}
    80  		}
    81  	}
    82  }
    83  
    84  func (s *CouchbaseSaver) saveItem(hostKey string, itemKey string, val IncomingValueData) error {
    85  	bucketKey := CouchbaseSaverBucketKey(hostKey, itemKey, val.Timestamp)
    86  
    87  	buf, err := msgpack.Marshal(val)
    88  	if err != nil {
    89  		return err
    90  	}
    91  
    92  	err = s.bucket.Do(bucketKey, func(mc *memcached.Client, vb uint16) (e error) {
    93  RETRY:
    94  		req := &gomemcached.MCRequest{
    95  			Opcode: gomemcached.APPEND,
    96  			VBucket: vb,
    97  			Key: []byte(bucketKey),
    98  			Cas: 0,
    99  			Opaque: 0,
   100  			Extras: nil,
   101  			Body: buf,
   102  		}
   103  
   104  		resp, e := mc.Send(req)
   105  		if resp == nil {
   106  			return
   107  		}
   108  
   109  		switch resp.Status {
   110  			case gomemcached.SUCCESS:
   111  				return
   112  			case gomemcached.NOT_STORED:
   113  				req := &gomemcached.MCRequest{
   114  					Opcode: gomemcached.ADD,
   115  					VBucket: vb,
   116  					Key: []byte(bucketKey),
   117  					Cas: 0,
   118  					Opaque: 0,
   119  					Extras: []byte{0, 0, 0, 0, 0, 0, 0, 0},
   120  					Body: buf,
   121  				}
   122  				binary.BigEndian.PutUint32(req.Extras[4:8], s.Ttl)
   123  
   124  				resp, e = mc.Send(req)
   125  				if resp == nil {
   126  					return
   127  				}
   128  
   129  				switch resp.Status {
   130  					case gomemcached.SUCCESS:
   131  						return
   132  					case gomemcached.KEY_EEXISTS:
   133  						goto RETRY
   134  					default:
   135  						return fmt.Errorf("ADD operation failed: %v", resp.Error())
   136  				}
   137  			default:
   138  				return fmt.Errorf("APPEND operation failed: %v", resp.Error())
   139  		}
   140  		panic("?!!")
   141  	});
   142  
   143  	return err
   144  }