code.gitea.io/gitea@v1.22.3/modules/queue/base_levelqueue_common.go (about)

     1  // Copyright 2023 The Gitea Authors. All rights reserved.
     2  // SPDX-License-Identifier: MIT
     3  
     4  package queue
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"path/filepath"
    10  	"strings"
    11  	"sync"
    12  	"time"
    13  
    14  	"code.gitea.io/gitea/modules/nosql"
    15  
    16  	"gitea.com/lunny/levelqueue"
    17  	"github.com/syndtr/goleveldb/leveldb"
    18  )
    19  
    20  // baseLevelQueuePushPoper is the common interface for levelqueue.Queue and levelqueue.UniqueQueue
    21  type baseLevelQueuePushPoper interface {
    22  	RPush(data []byte) error
    23  	LPop() ([]byte, error)
    24  	Len() int64
    25  }
    26  
    27  type baseLevelQueueCommonImpl struct {
    28  	length       int
    29  	internalFunc func() baseLevelQueuePushPoper
    30  	mu           *sync.Mutex
    31  }
    32  
    33  func (q *baseLevelQueueCommonImpl) PushItem(ctx context.Context, data []byte) error {
    34  	return backoffErr(ctx, backoffBegin, backoffUpper, time.After(pushBlockTime), func() (retry bool, err error) {
    35  		if q.mu != nil {
    36  			q.mu.Lock()
    37  			defer q.mu.Unlock()
    38  		}
    39  
    40  		cnt := int(q.internalFunc().Len())
    41  		if cnt >= q.length {
    42  			return true, nil
    43  		}
    44  		retry, err = false, q.internalFunc().RPush(data)
    45  		if err == levelqueue.ErrAlreadyInQueue {
    46  			err = ErrAlreadyInQueue
    47  		}
    48  		return retry, err
    49  	})
    50  }
    51  
    52  func (q *baseLevelQueueCommonImpl) PopItem(ctx context.Context) ([]byte, error) {
    53  	return backoffRetErr(ctx, backoffBegin, backoffUpper, infiniteTimerC, func() (retry bool, data []byte, err error) {
    54  		if q.mu != nil {
    55  			q.mu.Lock()
    56  			defer q.mu.Unlock()
    57  		}
    58  
    59  		data, err = q.internalFunc().LPop()
    60  		if err == levelqueue.ErrNotFound {
    61  			return true, nil, nil
    62  		}
    63  		if err != nil {
    64  			return false, nil, err
    65  		}
    66  		return false, data, nil
    67  	})
    68  }
    69  
    70  func baseLevelQueueCommon(cfg *BaseConfig, mu *sync.Mutex, internalFunc func() baseLevelQueuePushPoper) *baseLevelQueueCommonImpl {
    71  	return &baseLevelQueueCommonImpl{length: cfg.Length, mu: mu, internalFunc: internalFunc}
    72  }
    73  
    74  func prepareLevelDB(cfg *BaseConfig) (conn string, db *leveldb.DB, err error) {
    75  	if cfg.ConnStr == "" { // use data dir as conn str
    76  		if !filepath.IsAbs(cfg.DataFullDir) {
    77  			return "", nil, fmt.Errorf("invalid leveldb data dir (not absolute): %q", cfg.DataFullDir)
    78  		}
    79  		conn = cfg.DataFullDir
    80  	} else {
    81  		if !strings.HasPrefix(cfg.ConnStr, "leveldb://") {
    82  			return "", nil, fmt.Errorf("invalid leveldb connection string: %q", cfg.ConnStr)
    83  		}
    84  		conn = cfg.ConnStr
    85  	}
    86  	for i := 0; i < 10; i++ {
    87  		if db, err = nosql.GetManager().GetLevelDB(conn); err == nil {
    88  			break
    89  		}
    90  		time.Sleep(1 * time.Second)
    91  	}
    92  	return conn, db, err
    93  }