tikv_store.go 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368
  1. //go:build tikv
  2. // +build tikv
  3. package tikv
  4. import (
  5. "bytes"
  6. "context"
  7. "crypto/sha1"
  8. "fmt"
  9. "io"
  10. "strings"
  11. "github.com/seaweedfs/seaweedfs/weed/filer"
  12. "github.com/seaweedfs/seaweedfs/weed/glog"
  13. "github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
  14. "github.com/seaweedfs/seaweedfs/weed/util"
  15. "github.com/tikv/client-go/v2/config"
  16. "github.com/tikv/client-go/v2/txnkv"
  17. )
  18. var (
  19. _ filer.FilerStore = ((*TikvStore)(nil))
  20. )
  21. func init() {
  22. filer.Stores = append(filer.Stores, &TikvStore{})
  23. }
  24. type TikvStore struct {
  25. client *txnkv.Client
  26. deleteRangeConcurrency int
  27. onePC bool
  28. }
  29. // Basic APIs
  30. func (store *TikvStore) GetName() string {
  31. return "tikv"
  32. }
  33. func (store *TikvStore) Initialize(config util.Configuration, prefix string) error {
  34. ca := config.GetString(prefix + "ca_path")
  35. cert := config.GetString(prefix + "cert_path")
  36. key := config.GetString(prefix + "key_path")
  37. verify_cn := strings.Split(config.GetString(prefix+"verify_cn"), ",")
  38. pdAddrs := strings.Split(config.GetString(prefix+"pdaddrs"), ",")
  39. drc := config.GetInt(prefix + "deleterange_concurrency")
  40. if drc <= 0 {
  41. drc = 1
  42. }
  43. store.onePC = config.GetBool(prefix + "enable_1pc")
  44. store.deleteRangeConcurrency = drc
  45. return store.initialize(ca, cert, key, verify_cn, pdAddrs)
  46. }
  47. func (store *TikvStore) initialize(ca, cert, key string, verify_cn, pdAddrs []string) error {
  48. config.UpdateGlobal(func(conf *config.Config) {
  49. conf.Security = config.NewSecurity(ca, cert, key, verify_cn)
  50. })
  51. client, err := txnkv.NewClient(pdAddrs)
  52. store.client = client
  53. return err
  54. }
  55. func (store *TikvStore) Shutdown() {
  56. err := store.client.Close()
  57. if err != nil {
  58. glog.V(0).Infof("Shutdown TiKV client got error: %v", err)
  59. }
  60. }
  61. // ~ Basic APIs
  62. // Entry APIs
  63. func (store *TikvStore) InsertEntry(ctx context.Context, entry *filer.Entry) error {
  64. dir, name := entry.DirAndName()
  65. key := generateKey(dir, name)
  66. value, err := entry.EncodeAttributesAndChunks()
  67. if err != nil {
  68. return fmt.Errorf("encoding %s %+v: %v", entry.FullPath, entry.Attr, err)
  69. }
  70. txn, err := store.getTxn(ctx)
  71. if err != nil {
  72. return err
  73. }
  74. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  75. return txn.Set(key, value)
  76. })
  77. if err != nil {
  78. return fmt.Errorf("persisting %s : %v", entry.FullPath, err)
  79. }
  80. return nil
  81. }
  82. func (store *TikvStore) UpdateEntry(ctx context.Context, entry *filer.Entry) error {
  83. return store.InsertEntry(ctx, entry)
  84. }
  85. func (store *TikvStore) FindEntry(ctx context.Context, path util.FullPath) (*filer.Entry, error) {
  86. dir, name := path.DirAndName()
  87. key := generateKey(dir, name)
  88. txn, err := store.getTxn(ctx)
  89. if err != nil {
  90. return nil, err
  91. }
  92. var value []byte = nil
  93. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  94. val, err := txn.Get(context.TODO(), key)
  95. if err == nil {
  96. value = val
  97. }
  98. return err
  99. })
  100. if isNotExists(err) || value == nil {
  101. return nil, filer_pb.ErrNotFound
  102. }
  103. if err != nil {
  104. return nil, fmt.Errorf("get %s : %v", path, err)
  105. }
  106. entry := &filer.Entry{
  107. FullPath: path,
  108. }
  109. err = entry.DecodeAttributesAndChunks(value)
  110. if err != nil {
  111. return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
  112. }
  113. return entry, nil
  114. }
  115. func (store *TikvStore) DeleteEntry(ctx context.Context, path util.FullPath) error {
  116. dir, name := path.DirAndName()
  117. key := generateKey(dir, name)
  118. txn, err := store.getTxn(ctx)
  119. if err != nil {
  120. return err
  121. }
  122. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  123. return txn.Delete(key)
  124. })
  125. if err != nil {
  126. return fmt.Errorf("delete %s : %v", path, err)
  127. }
  128. return nil
  129. }
  130. // ~ Entry APIs
  131. // Directory APIs
  132. func (store *TikvStore) DeleteFolderChildren(ctx context.Context, path util.FullPath) error {
  133. directoryPrefix := genDirectoryKeyPrefix(path, "")
  134. txn, err := store.getTxn(ctx)
  135. if err != nil {
  136. return err
  137. }
  138. var (
  139. startKey []byte = nil
  140. endKey []byte = nil
  141. )
  142. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  143. iter, err := txn.Iter(directoryPrefix, nil)
  144. if err != nil {
  145. return err
  146. }
  147. defer iter.Close()
  148. for iter.Valid() {
  149. key := iter.Key()
  150. endKey = key
  151. if !bytes.HasPrefix(key, directoryPrefix) {
  152. break
  153. }
  154. if startKey == nil {
  155. startKey = key
  156. }
  157. err = iter.Next()
  158. if err != nil {
  159. return err
  160. }
  161. }
  162. // Only one Key matched just delete it.
  163. if startKey != nil && bytes.Equal(startKey, endKey) {
  164. return txn.Delete(startKey)
  165. }
  166. return nil
  167. })
  168. if err != nil {
  169. return fmt.Errorf("delete %s : %v", path, err)
  170. }
  171. if startKey != nil && endKey != nil && !bytes.Equal(startKey, endKey) {
  172. // has startKey and endKey and they are not equals, so use delete range
  173. _, err = store.client.DeleteRange(context.Background(), startKey, endKey, store.deleteRangeConcurrency)
  174. if err != nil {
  175. return fmt.Errorf("delete %s : %v", path, err)
  176. }
  177. }
  178. return err
  179. }
  180. func (store *TikvStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, eachEntryFunc filer.ListEachEntryFunc) (string, error) {
  181. return store.ListDirectoryPrefixedEntries(ctx, dirPath, startFileName, includeStartFile, limit, "", eachEntryFunc)
  182. }
  183. func (store *TikvStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (string, error) {
  184. lastFileName := ""
  185. directoryPrefix := genDirectoryKeyPrefix(dirPath, prefix)
  186. lastFileStart := directoryPrefix
  187. if startFileName != "" {
  188. lastFileStart = genDirectoryKeyPrefix(dirPath, startFileName)
  189. }
  190. txn, err := store.getTxn(ctx)
  191. if err != nil {
  192. return lastFileName, err
  193. }
  194. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  195. iter, err := txn.Iter(lastFileStart, nil)
  196. if err != nil {
  197. return err
  198. }
  199. defer iter.Close()
  200. for i := int64(0); i < limit && iter.Valid(); i++ {
  201. key := iter.Key()
  202. if !bytes.HasPrefix(key, directoryPrefix) {
  203. break
  204. }
  205. fileName := getNameFromKey(key)
  206. if fileName == "" || fileName == startFileName && !includeStartFile {
  207. if err := iter.Next(); err != nil {
  208. break
  209. } else {
  210. continue
  211. }
  212. }
  213. lastFileName = fileName
  214. entry := &filer.Entry{
  215. FullPath: util.NewFullPath(string(dirPath), fileName),
  216. }
  217. // println("list", entry.FullPath, "chunks", len(entry.GetChunks()))
  218. if decodeErr := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(iter.Value())); decodeErr != nil {
  219. err = decodeErr
  220. glog.V(0).Infof("list %s : %v", entry.FullPath, err)
  221. break
  222. }
  223. if err := iter.Next(); !eachEntryFunc(entry) || err != nil {
  224. break
  225. }
  226. }
  227. return nil
  228. })
  229. if err != nil {
  230. return lastFileName, fmt.Errorf("prefix list %s : %v", dirPath, err)
  231. }
  232. return lastFileName, nil
  233. }
  234. // ~ Directory APIs
  235. // Transaction Related APIs
  236. func (store *TikvStore) BeginTransaction(ctx context.Context) (context.Context, error) {
  237. tx, err := store.client.Begin()
  238. if err != nil {
  239. return ctx, err
  240. }
  241. if store.onePC {
  242. tx.SetEnable1PC(store.onePC)
  243. }
  244. return context.WithValue(ctx, "tx", tx), nil
  245. }
  246. func (store *TikvStore) CommitTransaction(ctx context.Context) error {
  247. if tx, ok := ctx.Value("tx").(*txnkv.KVTxn); ok {
  248. return tx.Commit(context.Background())
  249. }
  250. return nil
  251. }
  252. func (store *TikvStore) RollbackTransaction(ctx context.Context) error {
  253. if tx, ok := ctx.Value("tx").(*txnkv.KVTxn); ok {
  254. return tx.Rollback()
  255. }
  256. return nil
  257. }
  258. // ~ Transaction Related APIs
  259. // Transaction Wrapper
  260. type TxnWrapper struct {
  261. *txnkv.KVTxn
  262. inContext bool
  263. }
  264. func (w *TxnWrapper) RunInTxn(f func(txn *txnkv.KVTxn) error) error {
  265. err := f(w.KVTxn)
  266. if !w.inContext {
  267. if err != nil {
  268. w.KVTxn.Rollback()
  269. return err
  270. }
  271. w.KVTxn.Commit(context.Background())
  272. return nil
  273. }
  274. return err
  275. }
  276. func (store *TikvStore) getTxn(ctx context.Context) (*TxnWrapper, error) {
  277. if tx, ok := ctx.Value("tx").(*txnkv.KVTxn); ok {
  278. return &TxnWrapper{tx, true}, nil
  279. }
  280. txn, err := store.client.Begin()
  281. if err != nil {
  282. return nil, err
  283. }
  284. if store.onePC {
  285. txn.SetEnable1PC(store.onePC)
  286. }
  287. return &TxnWrapper{txn, false}, nil
  288. }
  289. // ~ Transaction Wrapper
  290. // Encoding Functions
  291. func hashToBytes(dir string) []byte {
  292. h := sha1.New()
  293. io.WriteString(h, dir)
  294. b := h.Sum(nil)
  295. return b
  296. }
  297. func generateKey(dirPath, fileName string) []byte {
  298. key := hashToBytes(dirPath)
  299. key = append(key, []byte(fileName)...)
  300. return key
  301. }
  302. func getNameFromKey(key []byte) string {
  303. return string(key[sha1.Size:])
  304. }
  305. func genDirectoryKeyPrefix(fullpath util.FullPath, startFileName string) (keyPrefix []byte) {
  306. keyPrefix = hashToBytes(string(fullpath))
  307. if len(startFileName) > 0 {
  308. keyPrefix = append(keyPrefix, []byte(startFileName)...)
  309. }
  310. return keyPrefix
  311. }
  312. func isNotExists(err error) bool {
  313. if err == nil {
  314. return false
  315. }
  316. if err.Error() == "not exist" {
  317. return true
  318. }
  319. return false
  320. }
  321. // ~ Encoding Functions