go.etcd.io/etcd@v3.3.27+incompatible/functional/runner/watch_command.go (about) 1 // Copyright 2016 The etcd Authors 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package runner 16 17 import ( 18 "context" 19 "errors" 20 "fmt" 21 "log" 22 "sync" 23 "time" 24 25 "github.com/coreos/etcd/clientv3" 26 "github.com/coreos/etcd/pkg/stringutil" 27 28 "github.com/spf13/cobra" 29 "golang.org/x/time/rate" 30 ) 31 32 var ( 33 runningTime time.Duration // time for which operation should be performed 34 noOfPrefixes int // total number of prefixes which will be watched upon 35 watchPerPrefix int // number of watchers per prefix 36 watchPrefix string // prefix append to keys in watcher 37 totalKeys int // total number of keys for operation 38 ) 39 40 // NewWatchCommand returns the cobra command for "watcher runner". 41 func NewWatchCommand() *cobra.Command { 42 cmd := &cobra.Command{ 43 Use: "watcher", 44 Short: "Performs watch operation", 45 Run: runWatcherFunc, 46 } 47 cmd.Flags().DurationVar(&runningTime, "running-time", 60, "number of seconds to run") 48 cmd.Flags().StringVar(&watchPrefix, "prefix", "", "the prefix to append on all keys") 49 cmd.Flags().IntVar(&noOfPrefixes, "total-prefixes", 10, "total no of prefixes to use") 50 cmd.Flags().IntVar(&watchPerPrefix, "watch-per-prefix", 10, "number of watchers per prefix") 51 cmd.Flags().IntVar(&totalKeys, "total-keys", 1000, "total number of keys to watch") 52 53 return cmd 54 } 55 56 func runWatcherFunc(cmd *cobra.Command, args []string) { 57 if len(args) > 0 { 58 ExitWithError(ExitBadArgs, errors.New("watcher does not take any argument")) 59 } 60 61 ctx := context.Background() 62 for round := 0; round < rounds || rounds <= 0; round++ { 63 fmt.Println("round", round) 64 performWatchOnPrefixes(ctx, cmd, round) 65 } 66 } 67 68 func performWatchOnPrefixes(ctx context.Context, cmd *cobra.Command, round int) { 69 keyPerPrefix := totalKeys / noOfPrefixes 70 prefixes := stringutil.UniqueStrings(5, noOfPrefixes) 71 keys := stringutil.RandomStrings(10, keyPerPrefix) 72 73 roundPrefix := fmt.Sprintf("%16x", round) 74 75 eps := endpointsFromFlag(cmd) 76 77 var ( 78 revision int64 79 wg sync.WaitGroup 80 gr *clientv3.GetResponse 81 err error 82 ) 83 84 client := newClient(eps, dialTimeout) 85 defer client.Close() 86 87 gr, err = getKey(ctx, client, "non-existent") 88 if err != nil { 89 log.Fatalf("failed to get the initial revision: %v", err) 90 } 91 revision = gr.Header.Revision 92 93 ctxt, cancel := context.WithDeadline(ctx, time.Now().Add(runningTime*time.Second)) 94 defer cancel() 95 96 // generate and put keys in cluster 97 limiter := rate.NewLimiter(rate.Limit(reqRate), reqRate) 98 99 go func() { 100 for _, key := range keys { 101 for _, prefix := range prefixes { 102 if err = limiter.Wait(ctxt); err != nil { 103 return 104 } 105 if err = putKeyAtMostOnce(ctxt, client, watchPrefix+"-"+roundPrefix+"-"+prefix+"-"+key); err != nil { 106 log.Fatalf("failed to put key: %v", err) 107 return 108 } 109 } 110 } 111 }() 112 113 ctxc, cancelc := context.WithCancel(ctx) 114 115 wcs := make([]clientv3.WatchChan, 0) 116 rcs := make([]*clientv3.Client, 0) 117 118 for _, prefix := range prefixes { 119 for j := 0; j < watchPerPrefix; j++ { 120 rc := newClient(eps, dialTimeout) 121 rcs = append(rcs, rc) 122 123 wprefix := watchPrefix + "-" + roundPrefix + "-" + prefix 124 125 wc := rc.Watch(ctxc, wprefix, clientv3.WithPrefix(), clientv3.WithRev(revision)) 126 wcs = append(wcs, wc) 127 128 wg.Add(1) 129 go func() { 130 defer wg.Done() 131 checkWatchResponse(wc, wprefix, keys) 132 }() 133 } 134 } 135 wg.Wait() 136 137 cancelc() 138 139 // verify all watch channels are closed 140 for e, wc := range wcs { 141 if _, ok := <-wc; ok { 142 log.Fatalf("expected wc to be closed, but received %v", e) 143 } 144 } 145 146 for _, rc := range rcs { 147 rc.Close() 148 } 149 150 if err = deletePrefix(ctx, client, watchPrefix); err != nil { 151 log.Fatalf("failed to clean up keys after test: %v", err) 152 } 153 } 154 155 func checkWatchResponse(wc clientv3.WatchChan, prefix string, keys []string) { 156 for n := 0; n < len(keys); { 157 wr, more := <-wc 158 if !more { 159 log.Fatalf("expect more keys (received %d/%d) for %s", n, len(keys), prefix) 160 } 161 for _, event := range wr.Events { 162 expectedKey := prefix + "-" + keys[n] 163 receivedKey := string(event.Kv.Key) 164 if expectedKey != receivedKey { 165 log.Fatalf("expected key %q, got %q for prefix : %q\n", expectedKey, receivedKey, prefix) 166 } 167 n++ 168 } 169 } 170 } 171 172 func putKeyAtMostOnce(ctx context.Context, client *clientv3.Client, key string) error { 173 gr, err := getKey(ctx, client, key) 174 if err != nil { 175 return err 176 } 177 178 var modrev int64 179 if len(gr.Kvs) > 0 { 180 modrev = gr.Kvs[0].ModRevision 181 } 182 183 for ctx.Err() == nil { 184 _, err := client.Txn(ctx).If(clientv3.Compare(clientv3.ModRevision(key), "=", modrev)).Then(clientv3.OpPut(key, key)).Commit() 185 186 if err == nil { 187 return nil 188 } 189 } 190 191 return ctx.Err() 192 } 193 194 func deletePrefix(ctx context.Context, client *clientv3.Client, key string) error { 195 for ctx.Err() == nil { 196 if _, err := client.Delete(ctx, key, clientv3.WithPrefix()); err == nil { 197 return nil 198 } 199 } 200 return ctx.Err() 201 } 202 203 func getKey(ctx context.Context, client *clientv3.Client, key string) (*clientv3.GetResponse, error) { 204 for ctx.Err() == nil { 205 if gr, err := client.Get(ctx, key); err == nil { 206 return gr, nil 207 } 208 } 209 return nil, ctx.Err() 210 }