github.com/polarismesh/polaris@v1.17.8/common/batchjob/future.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package batchjob
    19  
    20  import (
    21  	"context"
    22  	"errors"
    23  	"sync/atomic"
    24  	"time"
    25  )
    26  
    27  type Param interface{}
    28  
    29  type Future interface {
    30  	Param() Param
    31  	Done() (interface{}, error)
    32  	DoneTimeout(timeout time.Duration) (interface{}, error)
    33  	Cancel()
    34  	Reply(result interface{}, err error) error
    35  }
    36  
    37  type errorFuture struct {
    38  	task Param
    39  	err  error
    40  }
    41  
    42  func (f *errorFuture) Param() Param {
    43  	return f.task
    44  }
    45  
    46  func (f *errorFuture) Done() (interface{}, error) {
    47  	return nil, f.err
    48  }
    49  
    50  func (f *errorFuture) DoneTimeout(timeout time.Duration) (interface{}, error) {
    51  	return nil, f.err
    52  }
    53  
    54  func (f *errorFuture) Cancel() {
    55  }
    56  
    57  func (f *errorFuture) Reply(result interface{}, err error) error {
    58  	return nil
    59  }
    60  
    61  type future struct {
    62  	task      Param
    63  	setsignal chan struct{}
    64  	err       error
    65  	result    interface{}
    66  	replied   int32
    67  	ctx       context.Context
    68  	cancel    context.CancelFunc
    69  }
    70  
    71  func (f *future) Param() Param {
    72  	return f.task
    73  }
    74  
    75  func (f *future) Done() (interface{}, error) {
    76  	defer func() {
    77  		f.cancel()
    78  	}()
    79  	select {
    80  	case <-f.ctx.Done():
    81  		return nil, f.ctx.Err()
    82  	case <-f.setsignal:
    83  		return f.result, f.err
    84  	}
    85  }
    86  
    87  func (f *future) DoneTimeout(timeout time.Duration) (interface{}, error) {
    88  	timer := time.NewTimer(timeout)
    89  	defer func() {
    90  		timer.Stop()
    91  		f.cancel()
    92  	}()
    93  	select {
    94  	case <-timer.C:
    95  		return nil, context.DeadlineExceeded
    96  	case <-f.ctx.Done():
    97  		return nil, f.ctx.Err()
    98  	case <-f.setsignal:
    99  		f.cancel()
   100  		return f.result, f.err
   101  	}
   102  }
   103  
   104  func (f *future) Cancel() {
   105  	if atomic.CompareAndSwapInt32(&f.replied, 0, 1) {
   106  		close(f.setsignal)
   107  	}
   108  	f.cancel()
   109  }
   110  
   111  var (
   112  	ErrorReplyOnlyOnce = errors.New("reply only call once")
   113  	ErrorReplyCanceled = errors.New("reply canceled")
   114  )
   115  
   116  func (f *future) Reply(result interface{}, err error) error {
   117  	if !atomic.CompareAndSwapInt32(&f.replied, 0, 1) {
   118  		return ErrorReplyOnlyOnce
   119  	}
   120  	f.result = result
   121  	f.err = err
   122  	f.setsignal <- struct{}{}
   123  	close(f.setsignal)
   124  	return nil
   125  }