github.com/coyove/common@v0.0.0-20240403014525-f70e643f9de8/quadtree/redis/redis_adapter.go (about)

     1  package redis_adapter
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/gob"
     6  	"strconv"
     7  
     8  	"github.com/coyove/common/quadtree"
     9  	"github.com/gomodule/redigo/redis"
    10  )
    11  
    12  type Implement struct {
    13  	HGetAll func(id string) (map[string][]byte, error)
    14  	HSet    func(id, key string, v []byte) error
    15  	HSetNX  func(id, key string, v []byte) (existed bool, err error)
    16  	HDel    func(id, key string) error
    17  	Del     func(id string) error
    18  }
    19  
    20  type Database struct {
    21  	i Implement
    22  }
    23  
    24  func (db *Database) Load(id string) (quadtree.QuadTree, error) {
    25  	h, err := db.i.HGetAll(id)
    26  	if err != nil {
    27  		return quadtree.QuadTree{}, err
    28  	}
    29  	t := quadtree.QuadTree{}
    30  	if err := gob.NewDecoder(bytes.NewReader(h["t"])).Decode(&t); err != nil {
    31  		return quadtree.QuadTree{}, err
    32  	}
    33  	t.O[0], t.O[1], t.O[2], t.O[3] = string(h["0"]), string(h["1"]), string(h["2"]), string(h["3"])
    34  
    35  	el, err := db.i.HGetAll(id + "elems")
    36  	if err != nil {
    37  		return quadtree.QuadTree{}, err
    38  	}
    39  
    40  	t.Elems = map[quadtree.Point]quadtree.Element{}
    41  	for _, buf := range el {
    42  		var e quadtree.Element
    43  		if err := gob.NewDecoder(bytes.NewReader(buf)).Decode(&e); err != nil {
    44  			return quadtree.QuadTree{}, err
    45  		}
    46  		t.Elems[e.Point] = e
    47  	}
    48  	return t, nil
    49  }
    50  
    51  func (db *Database) Store(t quadtree.QuadTree) error {
    52  	buf := &bytes.Buffer{}
    53  	gob.NewEncoder(buf).Encode(t)
    54  	return db.i.HSet(t.ID, "t", buf.Bytes())
    55  }
    56  
    57  func (db *Database) StoreElement(id string, e quadtree.Element) error {
    58  	buf := &bytes.Buffer{}
    59  	gob.NewEncoder(buf).Encode(e)
    60  	return db.i.HSet(id+"elems", e.Point.Marshal(), buf.Bytes())
    61  }
    62  
    63  func (db *Database) DeleteAllElements(id string) error {
    64  	return db.i.Del(id + "elems")
    65  }
    66  
    67  func (db *Database) DeleteElement(id string, e quadtree.Element) error {
    68  	return db.i.HDel(id+"elems", e.Point.Marshal())
    69  }
    70  
    71  func (db *Database) StoreOrthant(id string, o int, oid string) (existed bool, err error) {
    72  	return db.i.HSetNX(id, strconv.Itoa(o), []byte(oid))
    73  }
    74  
    75  func New(redisInstance interface{}) *Database {
    76  	i, ok := redisInstance.(Implement)
    77  	if ok {
    78  		return &Database{i: i}
    79  	}
    80  	c, ok := redisInstance.(*redis.Pool)
    81  	if !ok {
    82  		c = redis.NewPool(func() (redis.Conn, error) {
    83  			return redis.Dial("tcp", redisInstance.(string))
    84  		}, 10)
    85  	}
    86  	return &Database{
    87  		i: Implement{
    88  			HGetAll: func(id string) (map[string][]byte, error) {
    89  				c := c.Get()
    90  				defer c.Close()
    91  				m2, err := redis.StringMap(c.Do("HGETALL", id))
    92  				if err != nil {
    93  					return nil, err
    94  				}
    95  				m := make(map[string][]byte, len(m2))
    96  				for k, v := range m2 {
    97  					m[k] = []byte(v)
    98  				}
    99  				return m, nil
   100  			},
   101  			HDel: func(id, key string) error {
   102  				c := c.Get()
   103  				defer c.Close()
   104  				_, err := redis.Bool(c.Do("HDEL", id, key))
   105  				return err
   106  			},
   107  			Del: func(id string) error {
   108  				c := c.Get()
   109  				defer c.Close()
   110  				_, err := redis.Bool(c.Do("DEL", id))
   111  				return err
   112  			},
   113  			HSet: func(id, key string, v []byte) error {
   114  				c := c.Get()
   115  				defer c.Close()
   116  				_, err := redis.Bool(c.Do("HSET", id, key, v))
   117  				return err
   118  			},
   119  			HSetNX: func(id, key string, v []byte) (bool, error) {
   120  				c := c.Get()
   121  				defer c.Close()
   122  				r, err := redis.Int(c.Do("HSETNX", id, key, v))
   123  				return r == 0, err
   124  			},
   125  		},
   126  	}
   127  }