github.com/aacfactory/fns@v1.2.86-0.20240310083819-80d667fc0a17/shareds/locker.go (about)

     1  /*
     2   * Copyright 2023 Wang Min Xiang
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   * 	http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   *
    16   */
    17  
    18  package shareds
    19  
    20  import (
    21  	"fmt"
    22  	"github.com/aacfactory/configures"
    23  	"github.com/aacfactory/errors"
    24  	"github.com/aacfactory/fns/commons/bytex"
    25  	"github.com/aacfactory/fns/commons/spinlock"
    26  	"github.com/aacfactory/fns/context"
    27  	"sync"
    28  	"sync/atomic"
    29  	"time"
    30  )
    31  
    32  var (
    33  	ErrLockTimeout = fmt.Errorf("fns: shared lockers lock timeout failed")
    34  )
    35  
    36  type Locker interface {
    37  	Lock(ctx context.Context) (err error)
    38  	Unlock(ctx context.Context) (err error)
    39  }
    40  
    41  type Lockers interface {
    42  	Acquire(ctx context.Context, key []byte, ttl time.Duration) (locker Locker, err error)
    43  	Close()
    44  }
    45  
    46  type LockersBuilder interface {
    47  	Build(ctx context.Context, config configures.Config) (lockers Lockers, err error)
    48  }
    49  
    50  type localLocker struct {
    51  	key       []byte
    52  	ttl       time.Duration
    53  	mutex     sync.Locker
    54  	releaseCh chan<- []byte
    55  	done      chan struct{}
    56  	locked    int64
    57  }
    58  
    59  func (locker *localLocker) Lock(ctx context.Context) (err error) {
    60  	deadline, hasDeadline := ctx.Deadline()
    61  	if hasDeadline && deadline.Before(time.Now().Add(locker.ttl)) {
    62  		locker.ttl = deadline.Sub(time.Now())
    63  	}
    64  	if locker.ttl == 0 {
    65  		atomic.StoreInt64(&locker.locked, 1)
    66  		locker.mutex.Lock()
    67  		return
    68  	}
    69  	go func(ctx context.Context, locker *localLocker) {
    70  		atomic.StoreInt64(&locker.locked, 1)
    71  		locker.mutex.Lock()
    72  		locker.done <- struct{}{}
    73  		close(locker.done)
    74  	}(ctx, locker)
    75  	select {
    76  	case <-locker.done:
    77  		break
    78  	case <-time.After(locker.ttl):
    79  		err = ErrLockTimeout
    80  		break
    81  	}
    82  	if err != nil {
    83  		_ = locker.Unlock(ctx)
    84  	}
    85  	return
    86  }
    87  
    88  func (locker *localLocker) Unlock(_ context.Context) (err error) {
    89  	if atomic.LoadInt64(&locker.locked) > 0 {
    90  		locker.mutex.Unlock()
    91  		atomic.StoreInt64(&locker.locked, 0)
    92  	}
    93  	locker.releaseCh <- locker.key
    94  	return
    95  }
    96  
    97  func LocalLockers() Lockers {
    98  	v := &localLockers{
    99  		mutex:     new(spinlock.Locker),
   100  		lockers:   make(map[string]*reuseLocker),
   101  		releaseCh: make(chan []byte, 10240),
   102  	}
   103  	go v.listenRelease()
   104  	return v
   105  }
   106  
   107  type localLockers struct {
   108  	mutex     sync.Locker
   109  	lockers   map[string]*reuseLocker
   110  	releaseCh chan []byte
   111  }
   112  
   113  func (lockers *localLockers) Acquire(_ context.Context, key []byte, ttl time.Duration) (locker Locker, err error) {
   114  	if key == nil || len(key) == 0 {
   115  		err = fmt.Errorf("%+v", errors.Warning("fns: shared lockers acquire failed").WithCause(errors.Warning("key is required")))
   116  		return
   117  	}
   118  	lockers.mutex.Lock()
   119  	rl, has := lockers.lockers[bytex.ToString(key)]
   120  	if !has {
   121  		rl = &reuseLocker{
   122  			mutex: new(spinlock.Locker),
   123  			times: 0,
   124  		}
   125  		lockers.lockers[bytex.ToString(key)] = rl
   126  	}
   127  	rl.times++
   128  	locker = &localLocker{
   129  		key:       key,
   130  		ttl:       ttl,
   131  		mutex:     rl.mutex,
   132  		releaseCh: lockers.releaseCh,
   133  		done:      make(chan struct{}, 1),
   134  		locked:    0,
   135  	}
   136  	lockers.mutex.Unlock()
   137  	return
   138  }
   139  
   140  func (lockers *localLockers) Close() {}
   141  
   142  func (lockers *localLockers) listenRelease() {
   143  	for {
   144  		key, ok := <-lockers.releaseCh
   145  		if !ok {
   146  			break
   147  		}
   148  		lockers.mutex.Lock()
   149  		v, has := lockers.lockers[bytex.ToString(key)]
   150  		if !has {
   151  			lockers.mutex.Unlock()
   152  			continue
   153  		}
   154  		v.times--
   155  		if v.times < 1 {
   156  			delete(lockers.lockers, bytex.ToString(key))
   157  		}
   158  		lockers.mutex.Unlock()
   159  	}
   160  	return
   161  }
   162  
   163  type reuseLocker struct {
   164  	mutex sync.Locker
   165  	times int64
   166  }