github.com/aacfactory/fns@v1.2.85/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/context"
    26  	"sync"
    27  	"sync/atomic"
    28  	"time"
    29  )
    30  
    31  var (
    32  	ErrLockTimeout = fmt.Errorf("fns: shared lockers lock timeout failed")
    33  )
    34  
    35  type Locker interface {
    36  	Lock(ctx context.Context) (err error)
    37  	Unlock(ctx context.Context) (err error)
    38  }
    39  
    40  type Lockers interface {
    41  	Acquire(ctx context.Context, key []byte, ttl time.Duration) (locker Locker, err error)
    42  	Close()
    43  }
    44  
    45  type LockersBuilder interface {
    46  	Build(ctx context.Context, config configures.Config) (lockers Lockers, err error)
    47  }
    48  
    49  type localLocker struct {
    50  	key       []byte
    51  	ttl       time.Duration
    52  	mutex     sync.Locker
    53  	releaseCh chan<- []byte
    54  	done      chan struct{}
    55  	locked    int64
    56  }
    57  
    58  func (locker *localLocker) Lock(ctx context.Context) (err error) {
    59  	deadline, hasDeadline := ctx.Deadline()
    60  	if hasDeadline && deadline.Before(time.Now().Add(locker.ttl)) {
    61  		locker.ttl = deadline.Sub(time.Now())
    62  	}
    63  	if locker.ttl == 0 {
    64  		atomic.StoreInt64(&locker.locked, 1)
    65  		locker.mutex.Lock()
    66  		return
    67  	}
    68  	go func(ctx context.Context, locker *localLocker) {
    69  		atomic.StoreInt64(&locker.locked, 1)
    70  		locker.mutex.Lock()
    71  		locker.done <- struct{}{}
    72  		close(locker.done)
    73  	}(ctx, locker)
    74  	select {
    75  	case <-locker.done:
    76  		break
    77  	case <-time.After(locker.ttl):
    78  		err = ErrLockTimeout
    79  		break
    80  	}
    81  	if err != nil {
    82  		_ = locker.Unlock(ctx)
    83  	}
    84  	return
    85  }
    86  
    87  func (locker *localLocker) Unlock(_ context.Context) (err error) {
    88  	if atomic.LoadInt64(&locker.locked) > 0 {
    89  		locker.mutex.Unlock()
    90  		atomic.StoreInt64(&locker.locked, 0)
    91  	}
    92  	locker.releaseCh <- locker.key
    93  	return
    94  }
    95  
    96  func LocalLockers() Lockers {
    97  	v := &localLockers{
    98  		mutex:     &sync.Mutex{},
    99  		lockers:   make(map[string]*reuseLocker),
   100  		releaseCh: make(chan []byte, 10240),
   101  	}
   102  	go v.listenRelease()
   103  	return v
   104  }
   105  
   106  type localLockers struct {
   107  	mutex     *sync.Mutex
   108  	lockers   map[string]*reuseLocker
   109  	releaseCh chan []byte
   110  }
   111  
   112  func (lockers *localLockers) Acquire(_ context.Context, key []byte, ttl time.Duration) (locker Locker, err error) {
   113  	if key == nil || len(key) == 0 {
   114  		err = fmt.Errorf("%+v", errors.Warning("fns: shared lockers acquire failed").WithCause(errors.Warning("key is required")))
   115  		return
   116  	}
   117  	lockers.mutex.Lock()
   118  	rl, has := lockers.lockers[bytex.ToString(key)]
   119  	if !has {
   120  		rl = &reuseLocker{
   121  			mutex: &sync.Mutex{},
   122  			times: 0,
   123  		}
   124  		lockers.lockers[bytex.ToString(key)] = rl
   125  	}
   126  	rl.times++
   127  	locker = &localLocker{
   128  		key:       key,
   129  		ttl:       ttl,
   130  		mutex:     rl.mutex,
   131  		releaseCh: lockers.releaseCh,
   132  		done:      make(chan struct{}, 1),
   133  		locked:    0,
   134  	}
   135  	lockers.mutex.Unlock()
   136  	return
   137  }
   138  
   139  func (lockers *localLockers) Close() {}
   140  
   141  func (lockers *localLockers) listenRelease() {
   142  	for {
   143  		key, ok := <-lockers.releaseCh
   144  		if !ok {
   145  			break
   146  		}
   147  		lockers.mutex.Lock()
   148  		v, has := lockers.lockers[bytex.ToString(key)]
   149  		if !has {
   150  			lockers.mutex.Unlock()
   151  			continue
   152  		}
   153  		v.times--
   154  		if v.times < 1 {
   155  			delete(lockers.lockers, bytex.ToString(key))
   156  		}
   157  		lockers.mutex.Unlock()
   158  	}
   159  	return
   160  }
   161  
   162  type reuseLocker struct {
   163  	mutex sync.Locker
   164  	times int64
   165  }