mcs/class/System.Threading.Tasks.Dataflow/System.Threading.Tasks.Dataflow/BatchedJoinBlock.cs

   1 // BatchedJoinBlock.cs
   2 //
   3 // Copyright (c) 2012 Petr Onderka
   4 //
   5 // Permission is hereby granted, free of charge, to any person obtaining a copy
   6 // of this software and associated documentation files (the "Software"), to deal
   7 // in the Software without restriction, including without limitation the rights
   8 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
   9 // copies of the Software, and to permit persons to whom the Software is
  10 // furnished to do so, subject to the following conditions:
  11 //
  12 // The above copyright notice and this permission notice shall be included in
  13 // all copies or substantial portions of the Software.
  14 //
  15 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  18 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  20 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  21 // THE SOFTWARE.
  22
  23 using System.Collections.Generic;
  24
  25 namespace System.Threading.Tasks.Dataflow {
  26         public sealed class BatchedJoinBlock<T1, T2> :
  27                 IReceivableSourceBlock<Tuple<IList<T1>, IList<T2>>> {
  28                 readonly GroupingDataflowBlockOptions options;
  29
  30                 readonly CompletionHelper completionHelper;
  31                 readonly OutgoingQueue<Tuple<IList<T1>, IList<T2>>> outgoing;
  32                 SpinLock batchLock;
  33
  34                 readonly JoinTarget<T1> target1;
  35                 readonly JoinTarget<T2> target2;
  36
  37                 int batchCount;
  38                 long numberOfGroups;
  39                 SpinLock batchCountLock;
  40
  41                 public BatchedJoinBlock (int batchSize)
  42                         : this (batchSize, GroupingDataflowBlockOptions.Default)
  43                 {
  44                 }
  45
  46                 public BatchedJoinBlock (int batchSize,
  47                                          GroupingDataflowBlockOptions dataflowBlockOptions)
  48                 {
  49                         if (batchSize <= 0)
  50                                 throw new ArgumentOutOfRangeException (
  51                                         "batchSize", batchSize, "The batchSize must be positive.");
  52                         if (dataflowBlockOptions == null)
  53                                 throw new ArgumentNullException ("dataflowBlockOptions");
  54                         if (!dataflowBlockOptions.Greedy)
  55                                 throw new ArgumentException (
  56                                         "Greedy must be true for this dataflow block.", "dataflowBlockOptions");
  57                         if (dataflowBlockOptions.BoundedCapacity != DataflowBlockOptions.Unbounded)
  58                                 throw new ArgumentException (
  59                                         "BoundedCapacity must be Unbounded or -1 for this dataflow block.",
  60                                         "dataflowBlockOptions");
  61
  62                         BatchSize = batchSize;
  63                         options = dataflowBlockOptions;
  64                         completionHelper = CompletionHelper.GetNew (dataflowBlockOptions);
  65
  66                         target1 = new JoinTarget<T1> (
  67                                 this, SignalTarget, completionHelper, () => outgoing.IsCompleted,
  68                                 dataflowBlockOptions, true, TryAdd);
  69                         target2 = new JoinTarget<T2> (
  70                                 this, SignalTarget, completionHelper, () => outgoing.IsCompleted,
  71                                 dataflowBlockOptions, true, TryAdd);
  72
  73                         outgoing = new OutgoingQueue<Tuple<IList<T1>, IList<T2>>> (
  74                                 this, completionHelper,
  75                                 () => target1.Buffer.IsCompleted || target2.Buffer.IsCompleted,
  76                                 _ =>
  77                                 {
  78                                         target1.DecreaseCount ();
  79                                         target2.DecreaseCount ();
  80                                 }, options);
  81                 }
  82
  83                 public int BatchSize { get; private set; }
  84
  85                 public ITargetBlock<T1> Target1 {
  86                         get { return target1; }
  87                 }
  88
  89                 public ITargetBlock<T2> Target2 {
  90                         get { return target2; }
  91                 }
  92
  93                 bool TryAdd ()
  94                 {
  95                         bool lockTaken = false;
  96                         try {
  97                                 batchCountLock.Enter (ref lockTaken);
  98
  99                                 if (options.MaxNumberOfGroups != -1
 100                                     && numberOfGroups + batchCount / BatchSize >= options.MaxNumberOfGroups)
 101                                         return false;
 102
 103                                 batchCount++;
 104                                 return true;
 105                         } finally {
 106                                 if (lockTaken)
 107                                         batchCountLock.Exit();
 108                         }
 109                 }
 110
 111                 void SignalTarget ()
 112                 {
 113                         bool lockTaken = false;
 114                         try {
 115                                 batchCountLock.Enter (ref lockTaken);
 116
 117                                 if (batchCount < BatchSize)
 118                                         return;
 119
 120                                 batchCount -= BatchSize;
 121                                 numberOfGroups++;
 122                         } finally {
 123                                 if (lockTaken)
 124                                         batchCountLock.Exit();
 125                         }
 126
 127                         MakeBatch (BatchSize);
 128                 }
 129
 130                 void MakeBatch (int batchSize)
 131                 {
 132                         if (batchSize == 0)
 133                                 return;
 134
 135                         var list1 = new List<T1> ();
 136                         var list2 = new List<T2> ();
 137
 138                         // lock is necessary here to make sure items are in the correct order
 139                         bool taken = false;
 140                         try {
 141                                 batchLock.Enter (ref taken);
 142
 143                                 int i = 0;
 144
 145                                 T1 item1;
 146                                 while (i < batchSize && target1.Buffer.TryTake (out item1)) {
 147                                         list1.Add (item1);
 148                                         i++;
 149                                 }
 150
 151                                 T2 item2;
 152                                 while (i < batchSize && target2.Buffer.TryTake (out item2)) {
 153                                         list2.Add (item2);
 154                                         i++;
 155                                 }
 156
 157                                 if (i < batchSize)
 158                                         throw new InvalidOperationException("Unexpected count of items.");
 159                         } finally {
 160                                 if (taken)
 161                                         batchLock.Exit ();
 162                         }
 163
 164                         var batch = Tuple.Create<IList<T1>, IList<T2>> (list1, list2);
 165
 166                         outgoing.AddData (batch);
 167
 168                         VerifyMaxNumberOfGroups ();
 169                 }
 170
 171                 void VerifyMaxNumberOfGroups ()
 172                 {
 173                         if (options.MaxNumberOfGroups == -1)
 174                                 return;
 175
 176                         bool shouldComplete;
 177
 178                         bool lockTaken = false;
 179                         try {
 180                                 batchCountLock.Enter (ref lockTaken);
 181
 182                                 shouldComplete = numberOfGroups >= options.MaxNumberOfGroups;
 183                         } finally {
 184                                 if (lockTaken)
 185                                         batchCountLock.Exit ();
 186                         }
 187
 188                         if (shouldComplete)
 189                                 Complete ();
 190                 }
 191
 192                 public Task Completion {
 193                         get { return completionHelper.Completion; }
 194                 }
 195
 196                 public void Complete ()
 197                 {
 198                         target1.Complete ();
 199                         target2.Complete ();
 200                         MakeBatch (batchCount);
 201                         outgoing.Complete ();
 202                 }
 203
 204                 void IDataflowBlock.Fault (Exception exception)
 205                 {
 206                         completionHelper.RequestFault (exception);
 207                 }
 208
 209                 Tuple<IList<T1>, IList<T2>> ISourceBlock<Tuple<IList<T1>, IList<T2>>>.ConsumeMessage (
 210                         DataflowMessageHeader messageHeader,
 211                         ITargetBlock<Tuple<IList<T1>, IList<T2>>> target,
 212                         out bool messageConsumed)
 213                 {
 214                         return outgoing.ConsumeMessage (messageHeader, target, out messageConsumed);
 215                 }
 216
 217                 public IDisposable LinkTo (ITargetBlock<Tuple<IList<T1>, IList<T2>>> target,
 218                                            DataflowLinkOptions linkOptions)
 219                 {
 220                         return outgoing.AddTarget(target, linkOptions);
 221                 }
 222
 223                 void ISourceBlock<Tuple<IList<T1>, IList<T2>>>.ReleaseReservation (
 224                         DataflowMessageHeader messageHeader,
 225                         ITargetBlock<Tuple<IList<T1>, IList<T2>>> target)
 226                 {
 227                         outgoing.ReleaseReservation (messageHeader, target);
 228                 }
 229
 230                 bool ISourceBlock<Tuple<IList<T1>, IList<T2>>>.ReserveMessage (
 231                         DataflowMessageHeader messageHeader,
 232                         ITargetBlock<Tuple<IList<T1>, IList<T2>>> target)
 233                 {
 234                         return outgoing.ReserveMessage (messageHeader, target);
 235                 }
 236
 237                 public bool TryReceive (Predicate<Tuple<IList<T1>, IList<T2>>> filter,
 238                                         out Tuple<IList<T1>, IList<T2>> item)
 239                 {
 240                         return outgoing.TryReceive (filter, out item);
 241                 }
 242
 243                 public bool TryReceiveAll (out IList<Tuple<IList<T1>, IList<T2>>> items)
 244                 {
 245                         return outgoing.TryReceiveAll (out items);
 246                 }
 247
 248                 public int OutputCount {
 249                         get { return outgoing.Count; }
 250                 }
 251
 252                 public override string ToString ()
 253                 {
 254                         return NameHelper.GetName (this, options);
 255                 }
 256         }
 257 }