github.com/aacfactory/fns@v1.2.86-0.20240310083819-80d667fc0a17/runtime/middleware.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 runtime
    19  
    20  import (
    21  	"github.com/aacfactory/errors"
    22  	"github.com/aacfactory/fns/commons/bytex"
    23  	"github.com/aacfactory/fns/commons/uid"
    24  	"github.com/aacfactory/fns/transports"
    25  	"github.com/aacfactory/logs"
    26  	"net/http"
    27  	"sync"
    28  )
    29  
    30  var (
    31  	ErrTooEarly    = errors.New(http.StatusTooEarly, "***TOO EARLY***", "fns: service is not ready, try later again")
    32  	ErrUnavailable = errors.Unavailable("fns: server is closed")
    33  )
    34  
    35  func Middleware(rt *Runtime) transports.Middleware {
    36  	return &middleware{
    37  		log:     nil,
    38  		rt:      rt,
    39  		counter: sync.WaitGroup{},
    40  	}
    41  }
    42  
    43  type middleware struct {
    44  	log     logs.Logger
    45  	rt      *Runtime
    46  	counter sync.WaitGroup
    47  }
    48  
    49  func (middle *middleware) Name() string {
    50  	return "runtime"
    51  }
    52  
    53  func (middle *middleware) Construct(options transports.MiddlewareOptions) error {
    54  	middle.log = options.Log
    55  	return nil
    56  }
    57  
    58  func (middle *middleware) Handler(next transports.Handler) transports.Handler {
    59  	return transports.HandlerFunc(func(w transports.ResponseWriter, r transports.Request) {
    60  		running, upped := middle.rt.Running()
    61  		if !running {
    62  			w.Header().Set(transports.ConnectionHeaderName, transports.CloseHeaderValue)
    63  			w.Failed(ErrUnavailable)
    64  			return
    65  		}
    66  		if !upped {
    67  			w.Header().Set(transports.ResponseRetryAfterHeaderName, bytex.FromString("3"))
    68  			w.Failed(ErrTooEarly)
    69  			return
    70  		}
    71  
    72  		middle.counter.Add(1)
    73  		// request Id
    74  		requestId := r.Header().Get(transports.RequestIdHeaderName)
    75  		if len(requestId) == 0 {
    76  			requestId = uid.Bytes()
    77  			r.Header().Set(transports.RequestIdHeaderName, requestId)
    78  		}
    79  		// set runtime into request context
    80  		With(r, middle.rt)
    81  		With(w, middle.rt)
    82  		// set request and response into context
    83  		transports.WithRequest(r, r)
    84  		transports.WithResponse(r, w)
    85  		// next
    86  		next.Handle(w, r)
    87  		// check hijacked
    88  		if w.Hijacked() {
    89  			middle.counter.Done()
    90  			return
    91  		}
    92  
    93  		// done
    94  		middle.counter.Done()
    95  	})
    96  }
    97  
    98  func (middle *middleware) Close() (err error) {
    99  	middle.counter.Wait()
   100  	return
   101  }