github.com/GuanceCloud/cliutils@v1.1.21/cache/cache.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  // Package cache wraps wal to handle local disk cache.
     7  package cache
     8  
     9  import (
    10  	"errors"
    11  	"os"
    12  	"sync"
    13  
    14  	"github.com/tidwall/wal"
    15  )
    16  
    17  type Cache struct {
    18  	l *wal.Log
    19  
    20  	lock *sync.Mutex
    21  	path string
    22  
    23  	totalMsg uint64
    24  
    25  	closed bool
    26  
    27  	cap, size uint64
    28  }
    29  
    30  func New(path string, capacity uint64) (*Cache, error) {
    31  	l, err := wal.Open(path, &wal.Options{NoCopy: true, LogFormat: wal.JSON})
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  
    36  	return &Cache{
    37  		l:    l,
    38  		path: path,
    39  		lock: &sync.Mutex{},
    40  		cap:  capacity,
    41  	}, nil
    42  }
    43  
    44  var (
    45  	ErrExceedCapacity = errors.New("exceed cache max capacity")
    46  	ErrCacheClosed    = errors.New("cache closed")
    47  )
    48  
    49  func (c *Cache) doPut(data []byte) error {
    50  	if c.closed {
    51  		if err := c.reopen(); err != nil {
    52  			return err
    53  		}
    54  	}
    55  
    56  	if c.cap > 0 && c.size+uint64(len(data)) > c.cap {
    57  		return ErrExceedCapacity
    58  	}
    59  
    60  	idx, err := c.l.LastIndex()
    61  	if err != nil {
    62  		return err
    63  	}
    64  
    65  	if err := c.l.Write(idx+1, data); err != nil {
    66  		return err
    67  	}
    68  
    69  	c.size += uint64(len(data))
    70  	c.totalMsg++
    71  	return nil
    72  }
    73  
    74  func (c *Cache) Put(data []byte) error {
    75  	c.lock.Lock()
    76  	defer c.lock.Unlock()
    77  	return c.doPut(data)
    78  }
    79  
    80  type Fn func([]byte) error
    81  
    82  func (c *Cache) Get(fn Fn) error {
    83  	c.lock.Lock()
    84  	defer c.lock.Unlock()
    85  
    86  	if c.closed {
    87  		return nil
    88  	}
    89  
    90  	idx, err := c.l.FirstIndex()
    91  	if err != nil {
    92  		return err
    93  	}
    94  
    95  	data, err := c.l.Read(idx)
    96  	if err != nil {
    97  		return err
    98  	}
    99  
   100  	if err := fn(data); err != nil {
   101  		return err
   102  	}
   103  
   104  	if c.lastEntry() {
   105  		if err := c.doClear(); err != nil {
   106  			return err
   107  		}
   108  
   109  		if err := c.reopen(); err != nil {
   110  			return err
   111  		}
   112  	} else {
   113  		if err := c.l.TruncateFront(idx + 1); err != nil {
   114  			return err
   115  		}
   116  	}
   117  	c.totalMsg--
   118  
   119  	return nil
   120  }
   121  
   122  func (c *Cache) doClear() error {
   123  	if err := c.l.Close(); err != nil {
   124  		return err
   125  	}
   126  
   127  	if err := os.RemoveAll(c.path); err != nil {
   128  		return err
   129  	}
   130  
   131  	c.size = 0
   132  	c.closed = true
   133  	return nil
   134  }
   135  
   136  func (c *Cache) Clear() error {
   137  	c.lock.Lock()
   138  	defer c.lock.Unlock()
   139  
   140  	return c.doClear()
   141  }
   142  
   143  func (c *Cache) Close() error {
   144  	c.lock.Lock()
   145  	defer c.lock.Unlock()
   146  
   147  	return c.l.Close()
   148  }
   149  
   150  func (c *Cache) lastEntry() bool {
   151  	first, err := c.l.FirstIndex()
   152  	if err != nil {
   153  		return false
   154  	}
   155  
   156  	last, err := c.l.LastIndex()
   157  	if err != nil {
   158  		return false
   159  	}
   160  
   161  	return first == last
   162  }
   163  
   164  func (c *Cache) reopen() error {
   165  	l, err := wal.Open(c.path, nil)
   166  	if err != nil {
   167  		return err
   168  	}
   169  	c.l = l
   170  	c.closed = false
   171  	return nil
   172  }