github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/indexcheckpoint/checkpoint.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package indexcheckpoint
    13  
    14  import (
    15  	"encoding/binary"
    16  	"fmt"
    17  	"os"
    18  	"path/filepath"
    19  
    20  	"github.com/pkg/errors"
    21  	"github.com/sirupsen/logrus"
    22  	bolt "go.etcd.io/bbolt"
    23  )
    24  
    25  var checkpointBucket = []byte("checkpoint")
    26  
    27  // Checkpoints keeps track of the last indexed vector id for each shard.
    28  // It stores the ids in a BoltDB file.
    29  type Checkpoints struct {
    30  	db   *bolt.DB
    31  	path string
    32  }
    33  
    34  func New(dir string, logger logrus.FieldLogger) (*Checkpoints, error) {
    35  	path := filepath.Join(dir, "index.db")
    36  
    37  	db, err := bolt.Open(path, 0o600, nil)
    38  	if err != nil {
    39  		return nil, errors.Wrapf(err, "open %q", path)
    40  	}
    41  
    42  	ic := Checkpoints{
    43  		db:   db,
    44  		path: path,
    45  	}
    46  
    47  	err = ic.initDB()
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  
    52  	return &ic, nil
    53  }
    54  
    55  func (c *Checkpoints) initDB() error {
    56  	err := c.db.Update(func(tx *bolt.Tx) error {
    57  		_, err := tx.CreateBucketIfNotExists(checkpointBucket)
    58  		return err
    59  	})
    60  
    61  	return errors.Wrap(err, "init db")
    62  }
    63  
    64  // Close the underlying DB
    65  func (c *Checkpoints) Close() {
    66  	c.db.Close()
    67  }
    68  
    69  func (c *Checkpoints) getID(shardID, targetVector string) string {
    70  	if targetVector != "" {
    71  		return fmt.Sprintf("%s_%s", shardID, targetVector)
    72  	}
    73  	return shardID
    74  }
    75  
    76  func (c *Checkpoints) Get(shardID, targetVector string) (count uint64, exists bool, err error) {
    77  	err = c.db.View(func(tx *bolt.Tx) error {
    78  		b := tx.Bucket(checkpointBucket)
    79  		v := b.Get([]byte(c.getID(shardID, targetVector)))
    80  		if v == nil {
    81  			return nil
    82  		}
    83  
    84  		count = binary.LittleEndian.Uint64(v)
    85  		exists = true
    86  		return nil
    87  	})
    88  	if err != nil {
    89  		return 0, false, errors.Wrap(err, "get checkpoint")
    90  	}
    91  
    92  	return
    93  }
    94  
    95  func (c *Checkpoints) Update(shardID, targetVector string, id uint64) error {
    96  	buf := make([]byte, 8)
    97  	binary.LittleEndian.PutUint64(buf, id)
    98  
    99  	err := c.db.Update(func(tx *bolt.Tx) error {
   100  		b := tx.Bucket(checkpointBucket)
   101  		return b.Put([]byte(c.getID(shardID, targetVector)), buf)
   102  	})
   103  	if err != nil {
   104  		return errors.Wrap(err, "update checkpoint")
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (c *Checkpoints) Delete(shardID, targetVector string) error {
   111  	err := c.db.Update(func(tx *bolt.Tx) error {
   112  		b := tx.Bucket(checkpointBucket)
   113  		return b.Delete([]byte(c.getID(shardID, targetVector)))
   114  	})
   115  	if err != nil {
   116  		return errors.Wrap(err, "delete checkpoint")
   117  	}
   118  
   119  	return nil
   120  }
   121  
   122  func (c *Checkpoints) Drop() error {
   123  	c.db.Close()
   124  	return os.Remove(c.Filename())
   125  }
   126  
   127  func (c *Checkpoints) Filename() string {
   128  	return c.path
   129  }