github.com/DARA-Project/GoDist-Scheduler@v0.0.0-20201030134746-668de4acea0d/examples/SimpleContext/context.go (about)

     1  package main
     2  
     3  import (
     4    "context"
     5    "fmt"
     6    "math/rand"
     7    "time"
     8  )
     9  
    10  //Slow function
    11  func sleepRandom(fromFunction string, ch chan int) {
    12    //defer cleanup
    13    defer func() { fmt.Println(fromFunction, "sleepRandom complete") }()
    14  
    15    //Perform a slow task
    16    //For illustration purpose,
    17    //Sleep here for random ms
    18    seed := time.Now().UnixNano()
    19    r := rand.New(rand.NewSource(seed))
    20    randomNumber := r.Intn(100)
    21    sleeptime := randomNumber + 100
    22    fmt.Println(fromFunction, "Starting sleep for", sleeptime, "ms")
    23    time.Sleep(time.Duration(sleeptime) * time.Millisecond)
    24    fmt.Println(fromFunction, "Waking up, slept for ", sleeptime, "ms")
    25  
    26    //write on the channel if it was passed in
    27    if ch != nil {
    28      ch <- sleeptime
    29    }
    30  }
    31  
    32  //Function that does slow processing with a context
    33  //Note that context is the first argument
    34  func sleepRandomContext(ctx context.Context, ch chan bool) {
    35  
    36    //Cleanup tasks
    37    //There are no contexts being created here
    38    //Hence, no canceling needed
    39    defer func() {
    40      fmt.Println("sleepRandomContext complete")
    41      ch <- true
    42    }()
    43  
    44    //Make a channel
    45    sleeptimeChan := make(chan int)
    46  
    47    //Start slow processing in a goroutine
    48    //Send a channel for communication
    49    go sleepRandom("sleepRandomContext", sleeptimeChan)
    50  
    51    //Use a select statement to exit out if context expires
    52    select {
    53    case <-ctx.Done():
    54      //If context is cancelled, this case is selected
    55      //This can happen if the timeout doWorkContext expires or
    56      //doWorkContext calls cancelFunction or main calls cancelFunction
    57      //Free up resources that may no longer be needed because of aborting the work
    58      //Signal all the goroutines that should stop work (use channels)
    59      //Usually, you would send something on channel, 
    60      //wait for goroutines to exit and then return
    61      //Or, use wait groups instead of channels for synchronization
    62      fmt.Println("sleepRandomContext: Time to return")
    63    case sleeptime := <-sleeptimeChan:
    64      //This case is selected when processing finishes before the context is cancelled
    65      fmt.Println("Slept for ", sleeptime, "ms")
    66    }
    67  }
    68  
    69  //A helper function, this can, in the real world do various things.
    70  //In this example, it is just calling one function.
    71  //Here, this could have just lived in main
    72  func doWorkContext(ctx context.Context) {
    73  
    74    //Derive a timeout context from context with cancel
    75    //Timeout in 150 ms
    76    //All the contexts derived from this will returns in 150 ms
    77    ctxWithTimeout, cancelFunction := context.WithTimeout(ctx, time.Duration(150)*time.Millisecond)
    78  
    79    //Cancel to release resources once the function is complete
    80    defer func() {
    81      fmt.Println("doWorkContext complete")
    82      cancelFunction()
    83    }()
    84  
    85    //Make channel and call context function
    86    //Can use wait groups as well for this particular case
    87    //As we do not use the return value sent on channel
    88    ch := make(chan bool)
    89    go sleepRandomContext(ctxWithTimeout, ch)
    90  
    91    //Use a select statement to exit out if context expires
    92    select {
    93    case <-ctx.Done():
    94      //This case is selected when the passed in context notifies to stop work
    95      //In this example, it will be notified when main calls cancelFunction
    96      fmt.Println("doWorkContext: Time to return")
    97    case <-ch:
    98      //This case is selected when processing finishes before the context is cancelled
    99      fmt.Println("sleepRandomContext returned")
   100    }
   101  }
   102  
   103  func main() {
   104    //Make a background context
   105    ctx := context.Background()
   106    //Derive a context with cancel
   107    ctxWithCancel, cancelFunction := context.WithCancel(ctx)
   108  
   109    //defer canceling so that all the resources are freed up 
   110    //For this and the derived contexts
   111    defer func() {
   112      fmt.Println("Main Defer: canceling context")
   113      cancelFunction()
   114    }()
   115  
   116    //Cancel context after a random time
   117    //This cancels the request after a random timeout
   118    //If this happens, all the contexts derived from this should return
   119    go func() {
   120      sleepRandom("Main", nil)
   121      cancelFunction()
   122      fmt.Println("Main Sleep complete. canceling context")
   123    }()
   124    //Do work
   125    doWorkContext(ctxWithCancel)
   126  }