1 /* 2 * Copyright (C) 2023 The Android Open Source Project 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 package com.android.federatedcompute.services.examplestore; 18 19 import static android.federatedcompute.common.ClientConstants.EXTRA_COLLECTION_NAME; 20 import static android.federatedcompute.common.ClientConstants.EXTRA_EXAMPLE_ITERATOR_RESULT; 21 import static android.federatedcompute.common.ClientConstants.STATUS_INTERNAL_ERROR; 22 23 import android.annotation.NonNull; 24 import android.federatedcompute.ExampleStoreIterator; 25 import android.federatedcompute.ExampleStoreIterator.IteratorCallback; 26 import android.federatedcompute.ExampleStoreService; 27 import android.federatedcompute.ExampleStoreService.QueryCallback; 28 import android.os.Bundle; 29 30 import com.google.common.collect.ImmutableList; 31 import com.google.protobuf.ByteString; 32 33 import org.tensorflow.example.BytesList; 34 import org.tensorflow.example.Example; 35 import org.tensorflow.example.Feature; 36 import org.tensorflow.example.Features; 37 38 import java.util.Iterator; 39 import java.util.List; 40 41 /** A sample ExampleStoreService implementation. */ 42 public class SampleExampleStoreService extends ExampleStoreService { 43 private static final String EXPECTED_COLLECTION_NAME = 44 "/federatedcompute.examplestoretest/test_collection"; 45 private static final Example EXAMPLE_PROTO_1 = 46 Example.newBuilder() 47 .setFeatures( 48 Features.newBuilder() 49 .putFeature( 50 "feature1", 51 Feature.newBuilder() 52 .setBytesList( 53 BytesList.newBuilder() 54 .addValue( 55 ByteString.copyFromUtf8( 56 "f1_value1"))) 57 .build())) 58 .build(); 59 60 @Override startQuery(@onNull Bundle params, @NonNull QueryCallback callback)61 public void startQuery(@NonNull Bundle params, @NonNull QueryCallback callback) { 62 String collection = params.getString(EXTRA_COLLECTION_NAME); 63 if (!collection.equals(EXPECTED_COLLECTION_NAME)) { 64 callback.onStartQueryFailure(STATUS_INTERNAL_ERROR); 65 return; 66 } 67 callback.onStartQuerySuccess( 68 new ListExampleStoreIterator(ImmutableList.of(EXAMPLE_PROTO_1))); 69 } 70 71 /** 72 * A simple {@link ExampleStoreIterator} that returns the contents of the {@link List} it's 73 * constructed with. 74 */ 75 private static class ListExampleStoreIterator implements ExampleStoreIterator { 76 private final Iterator<Example> mExampleIterator; 77 ListExampleStoreIterator(List<Example> examples)78 ListExampleStoreIterator(List<Example> examples) { 79 mExampleIterator = examples.iterator(); 80 } 81 82 @Override next(IteratorCallback callback)83 public synchronized void next(IteratorCallback callback) { 84 if (mExampleIterator.hasNext()) { 85 Bundle bundle = new Bundle(); 86 bundle.putByteArray( 87 EXTRA_EXAMPLE_ITERATOR_RESULT, mExampleIterator.next().toByteArray()); 88 callback.onIteratorNextSuccess(bundle); 89 } else { 90 callback.onIteratorNextSuccess(null); 91 } 92 } 93 94 @Override close()95 public void close() {} 96 } 97 } 98